3209. Number of Subarrays With AND Value of K
Description
Given an array of integers nums and an integer k, return the number of subarrays of nums where the bitwise AND of the elements of the subarray equals k.
Example 1:
Input: nums = [1,1,1], k = 1
Output: 6
Explanation:
All subarrays contain only 1's.
Example 2:
Input: nums = [1,1,2], k = 1
Output: 3
Explanation:
Subarrays having an AND value of 1 are: [1,1,2], [1,1,2], [1,1,2].
Example 3:
Input: nums = [1,2,3], k = 2
Output: 2
Explanation:
Subarrays having an AND value of 2 are: [1,2,3], [1,2,3].
Constraints:
1 <= nums.length <= 1050 <= nums[i], k <= 109
Solutions
Solution 1: Hash Table + Enumeration
According to the problem description, we need to find the result of the bitwise AND operation of elements from index $l$ to $r$ in the array $\textit{nums}$, that is, $\textit{nums}[l] \land \textit{nums}[l + 1] \land \cdots \land \textit{nums}[r]$, where $\land$ represents the bitwise AND operation.
If we fix the right endpoint $r$, then the range of the left endpoint $l$ is $[0, r]$. Since the sum of bitwise AND decreases monotonically as $l$ decreases, and the value of $nums[i]$ does not exceed $10^9$, the interval $[0, r]$ can have at most $30$ different values. Therefore, we can use a set to maintain all the values of $\textit{nums}[l] \land \textit{nums}[l + 1] \land \cdots \land \textit{nums}[r]$ and the number of times these values occur.
When we traverse from $r$ to $r+1$, the values with $r+1$ as the right endpoint are the values obtained by performing the bitwise AND operation of each value in the set with $nums[r + 1]$, plus $\textit{nums}[r + 1]$ itself.
Therefore, we only need to enumerate each value in the set and perform the bitwise AND operation with $\textit{nums[r]}$ to get all the values and their occurrences with $r$ as the right endpoint. Then, we need to add the occurrence count of $\textit{nums[r]}$. At this point, we add the occurrence count of value $k$ to the answer. Continue traversing $r$ until the entire array is traversed.
The time complexity is $O(n \times \log M)$, and the space complexity is $O(\log M)$. Here, $n$ and $M$ are the length of the array $\textit{nums}$ and the maximum value in the array $\textit{nums}$, respectively.
Python3
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
Java
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;
}
}
C++
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;
}
};
Go
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
}
TypeScript
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;
}