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

PythonServer Side ProgrammingProgramming

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
raja
Updated on 25-Oct-2021 06:41:10

Advertisements