Overview

This problem asks you to efficiently count the number of ways to select contiguous subsequences whose total weight is a multiple of \(K\), using prefix sums and properties of modular arithmetic.

Analysis

Naive Approach

If we consider trying all contiguous subsequences, there are \(O(N^2)\) combinations of starting point \(l\) and ending point \(r\). Furthermore, if we calculate the sum of each subsequence, the overall complexity becomes \(O(N^3)\), or \(O(N^2)\) even with prefix sums, which is too slow for \(N \leq 2 \times 10^5\).

Key Insight: Relationship Between Prefix Sums and Remainders

If we define the prefix sum \(S[i] = A_1 + A_2 + \cdots + A_i\) (with \(S[0] = 0\)), then the sum of the contiguous subsequence \([l+1, r]\) can be expressed as \(S[r] - S[l]\).

The condition for this sum to be a multiple of \(K\) is: $\(S[r] - S[l] \equiv 0 \pmod{K}\)$

Rearranging this gives: $\(S[r] \equiv S[l] \pmod{K}\)$

In other words, if we select two indices whose prefix sums have the same remainder when divided by \(K\), then the sum of the subsequence between them is a multiple of \(K\).

Concrete Example

For \(N = 4, K = 3, A = [1, 2, 3, 6]\):

\(i\)	\(S[i]\)	\(S[i] \mod 3\)
0	0	0
1	1	1
2	3	0
3	6	0
4	12	0

There are 4 indices with remainder 0: \(\{0, 2, 3, 4\}\), so choosing 2 from these contributes \(_4C_2 = 6\) combinations to the answer.

Algorithm

1. Calculate the prefix sum \(\mod K\) sequentially
2. Use a dictionary (hash map) to track how many times each remainder value has appeared
3. When calculating a new prefix sum remainder \(r\), add the number of times the same remainder \(r\) has already appeared to the answer
4. Then increment the count of remainder \(r\) by 1

Why does this count correctly?

As we look at each prefix sum \(S[i]\) in order, when \(S[i] \mod K = r\), if there are \(c\) previous prefix sums with the same remainder \(r\), we can pair \(S[i]\) with each of those \(c\) prefix sums. Each pair corresponds to one valid subsequence.

Setting count[0] = 1 in the initial state represents \(S[0] = 0\) (the state of selecting nothing).

Complexity

• Time complexity: \(O(N)\)
  • We traverse the array once, performing \(O(1)\) dictionary access for each element
• Space complexity: \(O(\min(N, K))\)
  • The number of distinct remainders stored in the dictionary is at most \(\min(N+1, K)\)

Implementation Notes

• Using defaultdict(int) automatically returns 0 for remainders that haven’t appeared yet
• Since \(K\) can be as large as \(10^9\), we need to use a dictionary instead of an array (an array would run out of memory)
• We only need to keep the remainder of the prefix sum divided by \(K\); there’s no need to store the prefix sums themselves in an array
• Be careful not to forget the initialization count[0] = 1, or you’ll miss counting subsequences that start from the beginning

Source Code

from collections import defaultdict

def solve():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    
    # Calculate prefix sum mod K
    # S[i] = (A[0] + A[1] + ... + A[i-1]) mod K
    # Sum of contiguous subsequence [l, r) is a multiple of K ⇔ S[r] ≡ S[l] (mod K)
    
    # Count occurrences of each prefix sum mod K
    count = defaultdict(int)
    count[0] = 1  # Empty prefix sum (before index 0)
    
    prefix_sum_mod = 0
    result = 0
    
    for i in range(N):
        prefix_sum_mod = (prefix_sum_mod + A[i]) % K
        # Add the count of previous prefix sums with the same mod K
        result += count[prefix_sum_mod]
        count[prefix_sum_mod] += 1
    
    print(result)

solve()

---

This editorial was generated by claude4.5opus.