跳转至

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;
}

评论