2588. Count the Number of Beautiful Subarrays

Description

You are given a 0-indexed integer array nums. In one operation, you can:

• Choose two different indices i and j such that 0 <= i, j < nums.length.
• Choose a non-negative integer k such that the kth bit (0-indexed) in the binary representation of nums[i] and nums[j] is 1.
• Subtract 2k from nums[i] and nums[j].

A subarray is beautiful if it is possible to make all of its elements equal to 0 after applying the above operation any number of times.

Return the number of beautiful subarrays in the array nums.

A subarray is a contiguous non-empty sequence of elements within an array.

 

Example 1:

Input: nums = [4,3,1,2,4]
Output: 2
Explanation: There are 2 beautiful subarrays in nums: [4,3,1,2,4] and [4,3,1,2,4].
- We can make all elements in the subarray [3,1,2] equal to 0 in the following way:
  - Choose [3, 1, 2] and k = 1. Subtract 21 from both numbers. The subarray becomes [1, 1, 0].
  - Choose [1, 1, 0] and k = 0. Subtract 20 from both numbers. The subarray becomes [0, 0, 0].
- We can make all elements in the subarray [4,3,1,2,4] equal to 0 in the following way:
  - Choose [4, 3, 1, 2, 4] and k = 2. Subtract 22 from both numbers. The subarray becomes [0, 3, 1, 2, 0].
  - Choose [0, 3, 1, 2, 0] and k = 0. Subtract 20 from both numbers. The subarray becomes [0, 2, 0, 2, 0].
  - Choose [0, 2, 0, 2, 0] and k = 1. Subtract 21 from both numbers. The subarray becomes [0, 0, 0, 0, 0].

Example 2:

Input: nums = [1,10,4]
Output: 0
Explanation: There are no beautiful subarrays in nums.

 

Constraints:

• 1 <= nums.length <= 105
• 0 <= nums[i] <= 106

Solutions

Solution 1: Prefix XOR + Hash Table

We observe that a subarray can become an array of all \(0\)s if and only if the number of \(1\)s on each binary bit of all elements in the subarray is even.

If there exist indices \(i\) and \(j\) such that \(i \lt j\) and the subarrays \(nums[0,..,i]\) and \(nums[0,..,j]\) have the same parity of the number of \(1\)s on each binary bit, then we can turn the subarray \(nums[i + 1,..,j]\) into an array of all \(0\)s.

Therefore, we can use the prefix XOR method and a hash table \(cnt\) to count the occurrences of each prefix XOR value. We traverse the array, for each element \(x\), we calculate its prefix XOR value \(mask\), then add the number of occurrences of \(mask\) to the answer. Then, we increase the number of occurrences of \(mask\) by \(1\).

Finally, we return the answer.

The time complexity is \(O(n)\), and the space complexity is \(O(n)\). Here, \(n\) is the length of the array \(nums\).

Python3JavaC++GoTypeScript

class Solution:
    def beautifulSubarrays(self, nums: List[int]) -> int:
        cnt = Counter({0: 1})
        ans = mask = 0
        for x in nums:
            mask ^= x
            ans += cnt[mask]
            cnt[mask] += 1
        return ans

class Solution {
    public long beautifulSubarrays(int[] nums) {
        Map<Integer, Integer> cnt = new HashMap<>();
        cnt.put(0, 1);
        long ans = 0;
        int mask = 0;
        for (int x : nums) {
            mask ^= x;
            ans += cnt.getOrDefault(mask, 0);
            cnt.merge(mask, 1, Integer::sum);
        }
        return ans;
    }
}

class Solution {
public:
    long long beautifulSubarrays(vector<int>& nums) {
        unordered_map<int, int> cnt{{0, 1}};
        long long ans = 0;
        int mask = 0;
        for (int x : nums) {
            mask ^= x;
            ans += cnt[mask];
            ++cnt[mask];
        }
        return ans;
    }
};

func beautifulSubarrays(nums []int) (ans int64) {
    cnt := map[int]int{0: 1}
    mask := 0
    for _, x := range nums {
        mask ^= x
        ans += int64(cnt[mask])
        cnt[mask]++
    }
    return
}

function beautifulSubarrays(nums: number[]): number {
    const cnt = new Map();
    cnt.set(0, 1);
    let ans = 0;
    let mask = 0;
    for (const x of nums) {
        mask ^= x;
        ans += cnt.get(mask) || 0;
        cnt.set(mask, (cnt.get(mask) || 0) + 1);
    }
    return ans;
}

Comments