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
Solution
Python3
class Solution:
def countSubarrays(self, nums: List[int], k: int) -> int:
N = len(nums)
M = max(nums).bit_length() + 1
def atLeast(k):
A = [0] * M
def calc(length):
v = 0
for b in range(M):
if A[b] == length:
v += (1 << b)
return v
i = 0
ans = 0
for j, x in enumerate(nums):
for b in range(M):
if x & (1 << b):
A[b] += 1
while i <= j and calc(j - i + 1) < k:
for b in range(M):
if nums[i] & (1 << b):
A[b] -= 1
i += 1
ans += (j - i + 1)
return ans
return atLeast(k) - atLeast(k + 1)C++
class Solution {
public:
long long countSubarrays(vector<int>& nums, int k) {
map<int, int> prev;
int n = nums.size();
long long cnt = 0;
for (int num : nums){
map<int, int> cur;
for (auto [key, v] : prev)
cur[key & num] += v;
cur[num]++;
cnt += cur[k];
prev = cur;
}
return cnt;
}
};