Prefix Sum Algorithm to Count Number of Nice Subarrays

Given an array of integers nums and an integer k. A subarray is called nice if there are k odd numbers on it. Return the number of nice sub-arrays.

Example 1:
Input: nums = [1,1,2,1,1], k = 3
Output: 2
Explanation: The only sub-arrays with 3 odd numbers are [1,1,2,1] and [1,2,1,1].

Example 2:
Input: nums = [2,4,6], k = 1
Output: 0
Explanation: There is no odd numbers in the array.

Example 3:
Input: nums = [2,2,2,1,2,2,1,2,2,2], k = 2
Output: 16

Constraints:
1 <= nums.length <= 50000
1 <= nums[i] <= 10^5
1 <= k <= nums.length

Hints:
After replacing each even by zero and every odd by one can we use prefix sum to find answer ?
Can we use two pointers to count number of sub-arrays ?
Can we store indices of odd numbers and for each k indices count number of sub-arrays contains them ?

Prefix Sum Algorithm to Count Sub Arrays

In order to apply the prefix sum algorithm, we can change the problem description to equivalent one by modifying the numbers such in a way that the even numbers are zero, and odd numbers are one.

Therefore, the K odd numbers are reflected by computing the sum of the subarray (as odd numbers are valued 1). And we compute and store the prefix sum in a hash map where the key is the sum and the value is count. Then we can sum up those corresponding values.

The initial value for hashmap with key (sum=0) should be set to 1.

The Java implementation based on HashMap. We can remove the check for key existence (e.g. containsKey) by using the getOrDefault method.

1
2
3
4
5
6
7
8
9
10
11
12
13

class Solution {
    public int numberOfSubarrays(int[] nums, int k) {
        int sum = 0, ans = 0;
        Map<Integer, Integer> prefix = new HashMap<>();
        prefix.put(0, 1);
        for (int i = 0; i < nums.length; ++ i) {
            sum += nums[i] & 1;
            ans += prefix.getOrDefault(sum - k, 0);
            prefix.put(sum, prefix.getOrDefault(sum, 0) + 1);
        }
        return ans;
    }
}

class Solution {
    public int numberOfSubarrays(int[] nums, int k) {
        int sum = 0, ans = 0;
        Map<Integer, Integer> prefix = new HashMap<>();
        prefix.put(0, 1);
        for (int i = 0; i < nums.length; ++ i) {
            sum += nums[i] & 1;
            ans += prefix.getOrDefault(sum - k, 0);
            prefix.put(sum, prefix.getOrDefault(sum, 0) + 1);
        }
        return ans;
    }
}

The C++ solution is based on unordered_map, and if the key isn’t existent, the value is the default value of the primitives.

1
2
3
4
5
6
7
8
9
10
11
12
13
14

class Solution {
public:
    int numberOfSubarrays(vector<int>& nums, int k) {
        unordered_map<int, int> prefix;
        int sum = 0, ans = 0;
        prefix[0] = 1;        
        for (int i = 0; i < nums.size(); ++ i) {
            sum += (nums[i] & 1);
            ans += prefix[sum - k];
            prefix[sum]++;
        }
        return ans;
    }
};

class Solution {
public:
    int numberOfSubarrays(vector<int>& nums, int k) {
        unordered_map<int, int> prefix;
        int sum = 0, ans = 0;
        prefix[0] = 1;        
        for (int i = 0; i < nums.size(); ++ i) {
            sum += (nums[i] & 1);
            ans += prefix[sum - k];
            prefix[sum]++;
        }
        return ans;
    }
};

Both implementations are complexity of O(N) time and O(N) space.

–EOF (The Ultimate Computing & Technology Blog) —