Program to find number of consecutive subsequences whose sum is divisible by k in Python

Suppose we have an array nums and a value k. We have to find number of consecutive subsequences whose sum is divisible by k.

So, if the input is like k = 3 nums = [1,2,3,4,1], then the output will be 4 because the subsequences are [3], [1,2], [1,2,3] and [2,3,4].

To solve this, we will follow these steps −

• x := An array of size k and fill with 0
• x[0] := 1
• r:= 0, s:= 0
• for each elem in nums, do
  • s :=(s + elem) mod k
  • r := r + x[s]
  • x[s] := x[s] + 1
• return r

Example

Let us see the following implementation to get better understanding −

def solve(k, nums):
   x = [0]*k
   x[0] = 1
   r=s=0
   for elem in nums:
      s = (s+elem) % k
      r += x[s]
      x[s] += 1
   return r

k = 3
nums = [1,2,3,4,1]
print(solve(k, nums))

Input

3, [1,2,3,4,1]

Output

4