LeetCode-in-Java

3729. Count Distinct Subarrays Divisible by K in Sorted Array

Hard

You are given an integer array nums sorted in non-descending order and a positive integer k.

A **non-empty subarrays** of nums is good if the sum of its elements is divisible by k.

Return an integer denoting the number of distinct good subarrays of nums.

Subarrays are distinct if their sequences of values are. For example, there are 3 distinct subarrays in [1, 1, 1], namely [1], [1, 1], and [1, 1, 1].

Example 1:

Input: nums = [1,2,3], k = 3

Output: 3

Explanation:

The good subarrays are [1, 2], [3], and [1, 2, 3]. For example, [1, 2, 3] is good because the sum of its elements is 1 + 2 + 3 = 6, and 6 % k = 6 % 3 = 0.

Example 2:

Input: nums = [2,2,2,2,2,2], k = 6

Output: 2

Explanation:

The good subarrays are [2, 2, 2] and [2, 2, 2, 2, 2, 2]. For example, [2, 2, 2] is good because the sum of its elements is 2 + 2 + 2 = 6, and 6 % k = 6 % 6 = 0.

Note that [2, 2, 2] is counted only once.

Constraints:

• 1 <= nums.length <= 105
• 1 <= nums[i] <= 109
• nums is sorted in non-descending order.
• 1 <= k <= 109

Solution

import java.util.HashMap;
import java.util.Map;

public class Solution {
    public long numGoodSubarrays(int[] nums, int k) {
        Map<Integer, Integer> cnt = new HashMap<>(nums.length, 1);
        cnt.put(0, 1);
        long sum = 0;
        int lastStart = 0;
        long ans = 0;
        for (int i = 0; i < nums.length; i++) {
            int x = nums[i];
            if (i > 0 && x != nums[i - 1]) {
                long s = sum;
                for (int t = i - lastStart; t > 0; t--) {
                    cnt.merge((int) (s % k), 1, Integer::sum);
                    s -= nums[i - 1];
                }
                lastStart = i;
            }
            sum += x;
            ans += cnt.getOrDefault((int) (sum % k), 0);
        }
        return ans;
    }
}

This site is open source. Improve this page.