LC 0974 [M] Subarray Sums Divisible by K - ALawliet/algorithms GitHub Wiki

After solving enough of these problems, it's not really a thing you need to come up with. When you hear "subarray sum", the first thing that comes to mind is prefix sum. Of course it might not help in every problem, but it's natural that it comes to mind. Once you think about prefix sums, you instantly interpret

"For how many subarrays is the sum divisible by K?" as

"subarrays is the sum" => "subarray sums" => pairs i<j of prefixsum[j]-prefixsum[i]

"For how many pairs i<j is prefixsum[j]-prefixsum[i] divisible by K?"

If you know your basic number theory, when you see "prefixsum[j]-prefixsum[i] is divisible by K", you might read that as "prefixsum[j] % K = prefixsum[i] % K" (assuming positive modulos)

So all in all, without really even trying to think, I just almost immediately read the problem as

"For how many pairs i<j is prefixsum[i] % K = prefixsum[j] % K?"

And that sounds like much more straightforward question than the original problem statement, so at that point it seems quite clear that thinking of prefix sums is the way to go.

class Solution:
    def subarraysDivByK(self, nums: List[int], k: int) -> int:
        prefixsum = 0
        count = 0
        d = {0:1} # prefix sum to freq, put a 0 in front to make it easier with edge case
        
        for j, x in enumerate(nums):
            prefixsum += x
            i = prefixsum % k
            
            if i in d:
                count += d[i]
                d[i] += 1
            else:
                d[i] = 1
        
        return count