D - お菓子の分配 / Distribution of Sweets Editorial by admin
GPT 5.2 HighOverview
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
- Let variable \(s\) be “the current prefix sum modulo \(K\)”, i.e., \(s = S_i \bmod K\) (initial value 0).
- Maintain a dictionary
cntthat stores “how many times each remainder has appeared so far”.- Initially, since \(S_0=0\) has appeared once, set
cnt[0]=1.
- Initially, since \(S_0=0\) has appeared once, set
- 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, soans += cnt[s] - Then update the occurrence count with
cnt[s] += 1
- 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] = 1will 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
% Kat 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: