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 need to find the number of consecutive subsequences whose sum is divisible by k.

So, if the input is like k = 3 and 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].

Algorithm

To solve this, we will follow these steps ?

• Create an array x of size k and fill with 0
• Set x[0] = 1 (to handle cases where prefix sum itself is divisible by k)
• Initialize result = 0 and sum = 0
• For each element in nums, do:
  • Update sum = (sum + element) mod k
  • Add x[sum] to result (count of previous occurrences)
  • Increment x[sum] by 1
• Return the result

How It Works

The algorithm uses the property that if two prefix sums have the same remainder when divided by k, then the subarray between them has a sum divisible by k. We track the frequency of each remainder in array x.

Example

Let us see the following implementation to get better understanding ?

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

k = 3
nums = [1, 2, 3, 4, 1]
print(f"Input: k = {k}, nums = {nums}")
print(f"Number of consecutive subsequences: {solve(k, nums)}")

Input: k = 3, nums = [1, 2, 3, 4, 1]
Number of consecutive subsequences: 4

Step-by-Step Trace

For k = 3 and nums = [1, 2, 3, 4, 1]:

• Initial: x = [1, 0, 0], result = 0, sum = 0
• Element 1: sum = 1, result += x[1] = 0, x = [1, 1, 0]
• Element 2: sum = 0, result += x[0] = 1, x = [2, 1, 0]
• Element 3: sum = 0, result += x[0] = 2, x = [3, 1, 0]
• Element 4: sum = 1, result += x[1] = 1, x = [3, 2, 0]
• Element 1: sum = 2, result += x[2] = 0, x = [3, 2, 1]

Final result: 4 subsequences with sum divisible by 3.

Conclusion

This algorithm efficiently finds consecutive subsequences with sum divisible by k using prefix sums and modular arithmetic. The time complexity is O(n) and space complexity is O(k).