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

• Leetcode 3624. Number of Integers With Popcount-Depth Equal to K II
  • 1. 解题思路
  • 2. 代码实现

• 题目链接：3624. Number of Integers With Popcount-Depth Equal to K II

1. 解题思路

这一题就是一个segment tree的变体，唯一的区别在于其操作不同于一般的简单求和或者求最大最小值，而是返回一个合并后的Counter即可。

而关于segment tree的相关内容，网络上实在是有很多了，我自己反正也有一篇博客《经典算法：Segment Tree》作为备忘，这里就不过多展开了，有兴趣的读者自己查查了解下即可。

2. 代码实现

给出python代码实现如下：

class SegmentTree:
    def __init__(self, arr):
        self.length = len(arr)
        self.tree = self.build(arr)

    def feature_func(self, *args):
        ret = defaultdict(int)
        for cnt in args:
            for k, v in cnt.items():
                ret[k] += v
        return ret

    def build(self, arr):
        n = len(arr)
        tree = [defaultdict(int) for _ in range(2*n)]
        for i in range(n):
            tree[i+n][arr[i]] += 1
        for i in range(n-1, 0, -1):
            tree[i] = self.feature_func(tree[i<<1], tree[(i<<1) | 1])
        return tree

    def update(self, idx, val):
        idx = idx + self.length
        self.tree[idx] = defaultdict(int)
        self.tree[idx][val] += 1
        while idx > 1:
            self.tree[idx>>1] = self.feature_func(self.tree[idx], self.tree[idx ^ 1])
            idx = idx>>1
        return

    def query(self, lb, rb):
        lb += self.length 
        rb += self.length
        nodes = []
        while lb < rb:
            if lb & 1 == 1:
                nodes.append(self.tree[lb])
                lb += 1
            if rb & 1 == 0:
                nodes.append(self.tree[rb])
                rb -= 1
            lb = lb >> 1
            rb = rb >> 1
        if lb == rb:
            nodes.append(self.tree[rb])
        return self.feature_func(*nodes)

class Solution:
    def popcountDepth(self, nums: List[int], queries: List[List[int]]) -> List[int]:
        
        def fn(num):
            if num == 1:
                return 0
            num = Counter(bin(num)[2:])["1"]
            return 1 + fn(num)

        nums = [fn(num) for num in nums]
        segment_tree = SegmentTree(nums)
        ans = []
        for query in queries:
            if len(query) == 3:
                _, idx, val = query
                segment_tree.update(idx, fn(val))
            else:
                _, l, r, k = query
                cnt = segment_tree.query(l, r)
                ans.append(cnt[k])
        return ans

提交代码评测得到：耗时6716ms，占用内存94.28MB。