leetcode weekly 119 - Subarray Sums Divisible by K

974. Subarray Sums Divisible by K

Given an array A of integers, return the number of (contiguous, non-empty) subarrays that have a sum divisible by K. 

Example 1:

Input: A = [4,5,0,-2,-3,1], K = 5
Output: 7
Explanation: There are 7 subarrays with a sum divisible by K = 5:
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]

Note:

1. 1 <= A.length <= 30000
2. -10000 <= A[i] <= 10000
3. 2 <= K <= 10000

---

Solution by neal_wu (Rank #3) in contest
Idea:

• See: https://leetcode.com/articles/subarray-sums-divisible-by-k/
• Calculate prefix sum of A, say P[i+1] = A[0] + ... + A[i], so subarray sum can be represented as P[j] - P[i] (j > i)
• Say C[i] = Count P[i] % K, i = 0..K-1
• Calculate combination sum for C[i] with i = 0..K-1

class Solution {
public:
    int subarraysDivByK(vector<int>& A, int K) {
        int n = A.size();
        vector<int> prefix_sum(n + 1, 0);

        for (int i = 0; i < n; i++)
            prefix_sum[i + 1] = ((prefix_sum[i] + A[i]) % K + K) % K;

Note! Because prefix_sum can be negative, so here need ((..)% K + K) % K;

• If prefix sum >= 0, the result is unchanged
• If prefix sum < 0, the result is complement value, e.g. P = -2, K = 3, ((P % K) + K) % K = 1

        vector<int> freq(K, 0);
        long long total = 0;

        for (int i = 0; i <= n; i++) {
            total += freq[prefix_sum[i]];
            freq[prefix_sum[i]]++;
        }

        return total;
    }
};

---

My second try after reading the answer. (I count positive and negative prefix sum separately)

  int subarraysDivByK(vector<int>& A, int K) {
    const int n = A.size();
    vector<int> prefsum(n + 1, 0);
    for (int i = 0; i < n; ++i)
      prefsum[i + 1] = prefsum[i] + A[i];
    
    vector<int> count(K, 0); // 0 ~ K-1
    vector<int> negCount(K+1, 0); // -K+1 ~ 0
    for (int p : prefsum) {
      int r = p % K;
      if (r < 0)
        ++negCount[abs(r)];
      else
        ++count[r];
    }
      
    int ans = 0;
    for (int i = 0; i < K; ++i) {
      int combine = count[i] + negCount[K - i];
      ans += combine * (combine - 1) / 2;
    }

    return ans;
  }