跳转至

3209. 子数组按位与值为 K 的数目

题目描述

给你一个整数数组 nums 和一个整数 k ，请你返回 nums 中有多少个子数组满足：子数组中所有元素按位 AND 的结果为 k 。

 

示例 1：

输入：nums = [1,1,1], k = 1

输出：6

解释：

所有子数组都只含有元素 1 。

示例 2：

输入：nums = [1,1,2], k = 1

输出：3

解释：

按位 AND 值为 1 的子数组包括：[1,1,2], [1,1,2], [1,1,2] 。

示例 3：

输入：nums = [1,2,3], k = 2

输出：2

解释：

按位 AND 值为 2 的子数组包括：[1,2,3], [1,2,3] 。

 

提示：

  • 1 <= nums.length <= 105
  • 0 <= nums[i], k <= 109

解法

方法一：哈希表 + 枚举

根据题目描述，我们需要求出数组 $\textit{nums}$ 下标 $l$ 到 $r$ 的元素的按位与运算的结果，即 $\textit{nums}[l] \land \textit{nums}[l + 1] \land \cdots \land \textit{nums}[r]$。其中 $\land$ 表示按位与运算。

如果我们每次固定右端点 $r$，那么左端点 $l$ 的范围是 $[0, r]$。由于按位与之和随着 $l$ 的减小而单调递减，并且 $nums[i]$ 的值不超过 $10^9$，因此区间 $[0, r]$ 最多只有 $30$ 种不同的值。因此，我们可以用一个集合来维护所有的 $\textit{nums}[l] \land \textit{nums}[l + 1] \land \cdots \land \textit{nums}[r]$ 的值，以及这些值出现的次数。

当我们从 $r$ 遍历到 $r+1$ 时，以 $r+1$ 为右端点的值，就是集合中每个值与 $nums[r + 1]$ 进行按位与运算得到的值，再加上 $\textit{nums}[r + 1]$ 本身。

因此，我们只需要枚举集合中的每个值，与 $\textit{nums[r]}$ 进行按位与运算，就可以得到以 $r$ 为右端点的所有值及其出现的次数。然后，我们还需要将 $\textit{nums[r]}$ 的出现次数加上去。此时，我们将值为 $k$ 的出现次数累加到答案中。继续遍历 $r$，直到遍历完整个数组。

时间复杂度 $O(n \times \log M)$，空间复杂度 $O(\log M)$。其中 $n$ 和 $M$ 分别是数组 $\textit{nums}$ 的长度和数组 $\textit{nums}$ 中元素的最大值。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
class Solution:
    def countSubarrays(self, nums: List[int], k: int) -> int:
        ans = 0
        pre = Counter()
        for x in nums:
            cur = Counter()
            for y, v in pre.items():
                cur[x & y] += v
            cur[x] += 1
            ans += cur[k]
            pre = cur
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution {
    public long countSubarrays(int[] nums, int k) {
        long ans = 0;
        Map<Integer, Integer> pre = new HashMap<>();
        for (int x : nums) {
            Map<Integer, Integer> cur = new HashMap<>();
            for (var e : pre.entrySet()) {
                int y = e.getKey(), v = e.getValue();
                cur.merge(x & y, v, Integer::sum);
            }
            cur.merge(x, 1, Integer::sum);
            ans += cur.getOrDefault(k, 0);
            pre = cur;
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
class Solution {
public:
    long long countSubarrays(vector<int>& nums, int k) {
        long long ans = 0;
        unordered_map<int, int> pre;
        for (int x : nums) {
            unordered_map<int, int> cur;
            for (auto& [y, v] : pre) {
                cur[x & y] += v;
            }
            cur[x]++;
            ans += cur[k];
            pre = cur;
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
func countSubarrays(nums []int, k int) (ans int64) {
    pre := map[int]int{}
    for _, x := range nums {
        cur := map[int]int{}
        for y, v := range pre {
            cur[x&y] += v
        }
        cur[x]++
        ans += int64(cur[k])
        pre = cur
    }
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
function countSubarrays(nums: number[], k: number): number {
    let ans = 0;
    let pre = new Map<number, number>();
    for (const x of nums) {
        const cur = new Map<number, number>();
        for (const [y, v] of pre) {
            const z = x & y;
            cur.set(z, (cur.get(z) || 0) + v);
        }
        cur.set(x, (cur.get(x) || 0) + 1);
        ans += cur.get(k) || 0;
        pre = cur;
    }
    return ans;
}

评论