Skip to content

3729. Count Distinct Subarrays Divisible by K in Sorted Array

Description

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

A subarray 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

Solutions

Solution 1

Thinking

The array is nondecreasing and \(n\le 10^5\). Prefix sums modulo \(k\) count every subarray whose sum is divisible by \(k\), but identical value sequences are over-counted. We then walk equal-value runs and subtract the extra copies whose length makes the run-sum divisible by \(k\), leaving distinct sequences.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution:
    def numGoodSubarrays(self, nums: List[int], k: int) -> int:
        cnt = Counter({0: 1})
        ans = s = 0
        for x in nums:
            s = (s + x) % k
            ans += cnt[s]
            cnt[s] += 1
        n = len(nums)
        i = 0
        while i < n:
            j = i + 1
            while j < n and nums[j] == nums[i]:
                j += 1
            m = j - i
            for h in range(1, m + 1):
                if (h * nums[i]) % k == 0:
                    ans -= m - h
            i = j
        return ans

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution {
    public long numGoodSubarrays(int[] nums, int k) {
        long ans = 0;
        int s = 0;
        Map<Integer, Integer> cnt = new HashMap<>();
        cnt.put(0, 1);
        for (int x : nums) {
            s = (s + x) % k;
            ans += cnt.getOrDefault(s, 0);
            cnt.merge(s, 1, Integer::sum);
        }
        int n = nums.length;
        for (int i = 0; i < n;) {
            int j = i + 1;
            while (j < n && nums[j] == nums[i]) {
                ++j;
            }
            int m = j - i;
            for (int h = 1; h <= m; ++h) {
                if (1L * nums[i] * h % k == 0) {
                    ans -= (m - h);
                }
            }
            i = j;
        }
        return ans;
    }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class Solution {
public:
    long long numGoodSubarrays(vector<int>& nums, int k) {
        long long ans = 0;
        int s = 0;
        unordered_map<int, int> cnt;
        cnt[0] = 1;
        for (int x : nums) {
            s = (s + x) % k;
            ans += cnt[s]++;
        }
        int n = nums.size();
        for (int i = 0; i < n;) {
            int j = i + 1;
            while (j < n && nums[j] == nums[i]) {
                ++j;
            }
            int m = j - i;
            for (int h = 1; h <= m; ++h) {
                if (1LL * nums[i] * h % k == 0) {
                    ans -= (m - h);
                }
            }
            i = j;
        }
        return ans;
    }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
func numGoodSubarrays(nums []int, k int) (ans int64) {
    s := 0
    cnt := map[int]int{0: 1}
    for _, x := range nums {
        s = (s + x) % k
        ans += int64(cnt[s])
        cnt[s]++
    }

    n := len(nums)
    for i := 0; i < n; {
        j := i + 1
        for j < n && nums[j] == nums[i] {
            j++
        }
        m := j - i
        for h := 1; h <= m; h++ {
            if int64(nums[i])*int64(h)%int64(k) == 0 {
                ans -= int64(m - h)
            }
        }
        i = j
    }
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
function numGoodSubarrays(nums: number[], k: number): number {
    let ans = 0;
    let s = 0;
    const cnt = new Map<number, number>();
    cnt.set(0, 1);

    for (const x of nums) {
        s = (s + x) % k;
        ans += cnt.get(s) ?? 0;
        cnt.set(s, (cnt.get(s) ?? 0) + 1);
    }

    const n = nums.length;
    for (let i = 0; i < n;) {
        let j = i + 1;
        while (j < n && nums[j] === nums[i]) ++j;
        const m = j - i;
        for (let h = 1; h <= m; ++h) {
            if ((nums[i] * h) % k === 0) {
                ans -= m - h;
            }
        }
        i = j;
    }

    return ans;
}

Comments