LeetCode Solution - 974. Subarray Sums Divisible by K

Problem Link: 974. Subarray Sums Divisible by K.

Function definition:

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var subarraysDivByK = function (nums, k) {}

Brute Force (TLE)

The first intuition for solving subarray problems is using prefix sum. For $n = nums.length$, define $prefixSum$ as $$\begin{aligned} prefixSum[-1] &= 0 \newline prefixSum[0] &= nums[0] \newline prefixSum[i] &= prefixSum[i-1] + nums[i], 1\leq i < n \end{aligned}$$

Then for $subarraySum = prefixSum[j]-prefixSum[i], -1\leq i<j<n$, we just need to count how many $subarraySum$s fulfill $subarraySum \bmod k = 0$. The time complexity is $O(n^2)$ and the space complexity is $O(n)$.

Count Remainders

Consider the requirement above $$(prefixSum[j]-prefixSum[i])\bmod k=0$$ We can make some transformations $$\begin{align} & (prefixSum[j]-prefixSum[i])\bmod k=0 \newline \Rightarrow & prefixSum[j] = prefixSum[i] + m\times k, m\in \mathbb{N} \newline \Rightarrow & (prefixSum[j])\bmod k = (prefixSum[i] + m\times k)\bmod k \newline \Rightarrow & (prefixSum[j])\bmod k = (prefixSum[i])\bmod k + (m\times k)\bmod k \newline \Rightarrow & (prefixSum[j])\bmod k = (prefixSum[i])\bmod k \newline \end{align}$$ So for $\sum_{a=i+1}^j nums[a]$, if $(prefixSum[j])\bmod k = (prefixSum[i])\bmod k$, then it fulfills the requirement. We define a map called remainderToCount to store the count of prefix sums’ remainders.

Once we have $r = remainderCount[remainder], remainder \neq 0$, we can pick $2$ from $r$ subarrays to get a required subarray. The number of subarrays we can get is $$C_r^2 = \frac{r!}{2!(r-2)!} = \frac{r(r-1)}{2}$$ A special case is $remainder = 0$. If $r=remainderCount[0]$, any one of the $r$ subarrays also fulfills the requirement. So we can pick $1$ or $2$ from $r$ subarrays to get a required subarray. The number of subarrays we can get is $$C_r^1 + C_r^2 = r+\frac{r!}{2!(r-2)!} = r+\frac{r(r-1)}{2}$$ Note that for some programming languages, $a\bmod b$ is not equivalent to $a\%b$ but equivalent to $((a\%b)+b)\%b$ as $-7\bmod 2=1$ and $-7 \% 2=-1$.

Implementation

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var subarraysDivByK = function (nums, k) {
    const N = nums.length;

    /** @type {Map<number, number>} */
    const remainderToCount = new Map();

    let currentPrefixSum = 0;

    let subArrayCount = 0;

    for (let i = 0; i < N; i++) {
        currentPrefixSum += nums[i];
        const reminder = ((currentPrefixSum % k) + k) % k;

        remainderToCount.set(
            reminder,
            (remainderToCount.get(reminder) ?? 0) + 1,
        );
    }

    for (const [, count] of remainderToCount) {
        // prevent overflow
        subArrayCount +=
            count % 2 === 1
                ? count * ((count - 1) / 2)
                : (count / 2) * (count - 1);
    }

    // spacial case of remainder = 0
    subArrayCount += remainderToCount.get(0) ?? 0;

    return subArrayCount;
};

The popular solution on LeetCode is equivalent to the solution above. It calculates subArrayCount while calculating currentPrefixSum and handles the special case before the calculation.

/**
 * @param {number[]} nums
 * @param {number} k
 * @return {number}
 */
var subarraysDivByK = function (nums, k) {
    const N = nums.length;

    /** @type {Map<number, number>} */
    const remainderToCount = new Map();

    // Every time we get a subarray fulfills currentPrefixSum % k === 0, 
    // we additionally add the current subarray to subArrayCount
    remainderToCount.set(0, 1);    

    let currentPrefixSum = 0;

    let subArrayCount = 0;

    for (let i = 0; i < N; i++) {
        currentPrefixSum += nums[i];
        const reminder = ((currentPrefixSum % k) + k) % k;

        subArrayCount += remainderToCount.get(reminder) ?? 0;

        remainderToCount.set(
            reminder,
            (remainderToCount.get(reminder) ?? 0) + 1,
        );
    }

    return subArrayCount;
};