Medium

Subarray Sums Divisible by K

LAB

Return the count of contiguous subarrays whose sum is divisible by k.

EXAMPLES

Example 1

Input
{
  "nums": [
    4,
    5,
    0,
    -2,
    -3,
    1
  ],
  "k": 5
}

Output
7

FUNCTION SHAPE

nums: intArrayk: int→int

SOLUTION NOTE

This is the exact same idea - we track prefix sums modulo k. If two prefix sums have the same remainder, their difference is divisible by k.

Reveal reference solution +

pythonREFERENCE

def subarraysDivByK(self, nums: List[int], k: int) -> int:
    # [ pref ][          0    ]
    # [    pref     ][   0    ]
    # [           pref        ]

    # For all indices j < i where pref[j] == pref[i]:
    # the subarray [j+1,i] is divisible by k.

    pref_mods = Counter({0:1})  # pref_mod -> freq
    pref, res = 0, 0
    for i in range(len(nums)):
        pref = (pref + nums[i]) % k
        res += pref_mods[pref]
        pref_mods[pref] += 1
    return res

TimeO(n)

SpaceO(k)

Open on LeetCode ↗