Official

D - お菓子の分配 / Distribution of Sweets Editorial by admin

GPT 5.2 High

Overview

When selecting a contiguous interval (contiguous subsequence) of sweets, count the number of intervals whose total weight is a multiple of \(K\).

Approach

The sum of a contiguous subsequence can be handled efficiently using prefix sums. If we define the prefix sum as \(S_i = A_1 + A_2 + \cdots + A_i\) (with \(S_0=0\)), then the sum of the interval \([l, r]\) is:

  • \(A_l + \cdots + A_r = S_r - S_{l-1}\)

The condition for this to be a multiple of \(K\) is:

  • \(S_r - S_{l-1} \equiv 0 \pmod K\)
  • In other words, \(S_r \equiv S_{l-1} \pmod K\)

Therefore, the answer is the number of pairs of positions that have the same “prefix sum modulo \(K\)”.


Naively checking all intervals would require examining \(O(N^2)\) intervals (up to about \(4\times 10^{10}\) when \(N\) is around \(2\times 10^5\)), which won’t finish within the time limit (TLE).

Instead, we manage the count of “prefix sum remainders seen so far” while scanning from left to right, completing the counting in a single pass.

Concrete Example

Suppose \(K=5\) and the remainder sequence (as \(S_i \bmod 5\)) is: \([S_0,S_1,S_2,S_3] \bmod 5 = [0, 2, 0, 2]\).

  • Remainder \(0\) appears 2 times → the number of pairs with the same remainder is \(\binom{2}{2}=1\)
  • Remainder \(2\) also appears 2 times → \(\binom{2}{2}=1\)

In total, 2 interval sums are multiples of \(5\). (In the implementation, we achieve the same thing by incrementally adding “how many times the same remainder as the current one has appeared in the past”.)

Algorithm

  1. Let variable \(s\) be “the current prefix sum modulo \(K\)”, i.e., \(s = S_i \bmod K\) (initial value 0).
  2. Maintain a dictionary cnt that stores “how many times each remainder has appeared so far”.
    • Initially, since \(S_0=0\) has appeared once, set cnt[0]=1.
  3. Process each \(A_i\) from left to right:
    • \(s \leftarrow (s + A_i) \bmod K\)
    • If the current remainder \(s\) has appeared cnt[s] times in the past, that many “intervals satisfying the condition” are added, so ans += cnt[s]
    • Then update the occurrence count with cnt[s] += 1
  4. Output ans.

This method counts all pairs where \(S_r \equiv S_{l-1} \pmod K\) without duplication.

Complexity

  • Time complexity: \(O(N)\) (a single pass, with dictionary operations taking amortized constant time)
  • Space complexity: \(O(M)\) (where \(M \le \min(N, K)\) is the number of distinct remainders that appear. Since \(K\) can be large, we use a dictionary instead of an array)

Implementation Notes

  • Forgetting cnt[0] = 1 will cause you to miss intervals starting from the beginning (\(l=1\)).

  • Since \(K\) can be up to \(10^9\), it is impossible to store remainder counts in an array of length \(K\). Use a dictionary such as defaultdict(int) to manage them.

  • If you keep adding to the prefix sum directly, it becomes very large, so take % K at each step to only track the remainder (in Python there is no overflow, but this is advantageous in terms of computational efficiency and ease of handling).

    Source Code

import sys
from collections import defaultdict

def main():
    input = sys.stdin.buffer.readline
    N, K = map(int, input().split())
    A = list(map(int, input().split()))

    cnt = defaultdict(int)
    cnt[0] = 1
    s = 0
    ans = 0

    for x in A:
        s = (s + x) % K
        ans += cnt[s]
        cnt[s] += 1

    print(ans)

if __name__ == "__main__":
    main()

This editorial was generated by gpt-5.2-high.

posted:
last update: