LeetCode 3624 - Number of Integers With Popcount-Depth Equal to K II

LeetCode Problem 3624

Difficulty: 🔴 Hard
Topics: Array, Divide and Conquer, Binary Indexed Tree, Segment Tree

Solution

Problem Understanding

The problem asks us to efficiently handle queries on an array nums involving a special property called the popcount-depth. The popcount-depth of a number x is defined as the number of times we repeatedly apply the binary population count (the number of 1 bits in x) until we reach 1. For example, for x = 7, the sequence is 7 → 3 → 2 → 1, giving a popcount-depth of 3.

We have two types of queries. The first type [1, l, r, k] asks how many numbers in the subarray nums[l:r+1] have a popcount-depth equal to k. The second type [2, idx, val] updates nums[idx] to a new value val. The constraints indicate that both nums and queries can have up to 10^5 elements, and values in nums can be as large as 10^15. This rules out naive approaches because recomputing popcount-depth for each query would be too slow. The problem guarantees valid indices and small values of k (0 to 5), which is a key optimization opportunity.

Edge cases include extremely large numbers that require multiple popcount iterations, updates that change a number’s depth significantly, and queries that span the entire array.

Approaches

The brute-force approach would be to recompute the popcount-depth for each query of type 1 by iterating over the subarray and computing the depth for each number. Updates would simply replace the number in nums. While correct, this approach is too slow because it could require O(n) computation per query, leading to O(n * q) overall time, which is unacceptable for n, q up to 10^5.

The optimal approach leverages two observations. First, popcount-depth values are very small (0 to 5) because repeated popcount reduces numbers quickly. Second, range queries with updates suggest using a Segment Tree (or Binary Indexed Tree) to maintain counts of numbers by depth. Each node in the segment tree stores an array of counts for depths 0 to 5. Updates propagate changes in depth counts, and range queries sum the relevant counts efficiently.

Algorithm Walkthrough

1. Precompute popcount-depths: Write a helper function to compute the popcount-depth of any number x. Repeatedly apply popcount until 1 is reached. This ensures updates and queries can quickly determine a number's depth.
2. Build Segment Tree: Each node of the segment tree stores an array of 6 integers, counts[0..5], representing how many numbers in that node’s range have popcount-depth equal to each value.
3. Segment Tree Build: For leaf nodes, initialize counts based on the depth of the single number. For internal nodes, sum counts from the left and right children.
4. Query Operation: To answer [1, l, r, k], traverse the segment tree to sum counts for depth k over the subarray [l, r]. This uses standard segment tree range query logic.
5. Update Operation: To handle [2, idx, val], compute the new depth for val, update the leaf node at idx by decrementing the old depth count and incrementing the new depth count, and propagate the changes up the tree.
6. Iterate Queries: For each query, either perform a query or update as described. Collect results from type 1 queries into the answer list.

Why it works: The segment tree guarantees that each query or update takes O(log n) time. Because counts are stored explicitly for each depth, range queries for depth k are direct sums, and updates propagate correctly. This ensures correctness and efficiency.

Python Solution

from typing import List

class Solution:
    def popcountDepth(self, nums: List[int], queries: List[List[int]]) -> List[int]:
        def popcount_depth(x: int) -> int:
            depth = 0
            while x != 1:
                x = bin(x).count('1')
                depth += 1
            return depth

        n = len(nums)
        max_depth = 5

        # Segment Tree Node representation
        tree = [[0]*(max_depth+1) for _ in range(4*n)]
        depths = [popcount_depth(x) for x in nums]

        def build(node: int, l: int, r: int):
            if l == r:
                tree[node][depths[l]] = 1
                return
            mid = (l + r) // 2
            build(node*2, l, mid)
            build(node*2+1, mid+1, r)
            for i in range(max_depth+1):
                tree[node][i] = tree[node*2][i] + tree[node*2+1][i]

        def update(node: int, l: int, r: int, idx: int, old_d: int, new_d: int):
            if l == r:
                tree[node][old_d] -= 1
                tree[node][new_d] += 1
                return
            mid = (l + r) // 2
            if idx <= mid:
                update(node*2, l, mid, idx, old_d, new_d)
            else:
                update(node*2+1, mid+1, r, idx, old_d, new_d)
            for i in range(max_depth+1):
                tree[node][i] = tree[node*2][i] + tree[node*2+1][i]

        def query(node: int, l: int, r: int, ql: int, qr: int, k: int) -> int:
            if qr < l or r < ql:
                return 0
            if ql <= l and r <= qr:
                return tree[node][k]
            mid = (l + r) // 2
            return query(node*2, l, mid, ql, qr, k) + query(node*2+1, mid+1, r, ql, qr, k)

        build(1, 0, n-1)
        answer = []
        for q in queries:
            if q[0] == 1:
                _, l, r, k = q
                answer.append(query(1, 0, n-1, l, r, k))
            else:
                _, idx, val = q
                old_depth = depths[idx]
                new_depth = popcount_depth(val)
                depths[idx] = new_depth
                update(1, 0, n-1, idx, old_depth, new_depth)

        return answer

Explanation: We precompute the initial popcount-depths and build a segment tree where each node stores counts for depths 0 to 5. Queries sum counts for a depth, and updates modify the counts and propagate changes. This keeps both operations O(log n).

Go Solution

package main

func popcountDepth(nums []int64, queries [][]int64) []int {
    n := len(nums)
    maxDepth := 5
    tree := make([][6]int, 4*n)
    depths := make([]int, n)

    var popcountDepth func(int64) int
    popcountDepth = func(x int64) int {
        depth := 0
        for x != 1 {
            cnt := 0
            y := x
            for y > 0 {
                cnt += int(y & 1)
                y >>= 1
            }
            x = int64(cnt)
            depth++
        }
        return depth
    }

    for i, val := range nums {
        depths[i] = popcountDepth(val)
    }

    var build func(int, int, int)
    build = func(node, l, r int) {
        if l == r {
            tree[node][depths[l]] = 1
            return
        }
        mid := (l + r) / 2
        build(node*2, l, mid)
        build(node*2+1, mid+1, r)
        for i := 0; i <= maxDepth; i++ {
            tree[node][i] = tree[node*2][i] + tree[node*2+1][i]
        }
    }

    var update func(int, int, int, int, int, int)
    update = func(node, l, r, idx, oldD, newD int) {
        if l == r {
            tree[node][oldD]--
            tree[node][newD]++
            return
        }
        mid := (l + r) / 2
        if idx <= mid {
            update(node*2, l, mid, idx, oldD, newD)
        } else {
            update(node*2+1, mid+1, r, idx, oldD, newD)
        }
        for i := 0; i <= maxDepth; i++ {
            tree[node][i] = tree[node*2][i] + tree[node*2+