D - お菓子の分配 / Distribution of Sweets Editorial by admin
Gemini 3.0 FlashOverview
Given a sequence of \(N\) elements, find the total number of contiguous subsequences whose sum is a multiple of \(K\).
Analysis
Naive Approach
The number of ways to choose a contiguous subsequence is \(O(N^2)\), where we pick the left endpoint \(i\) and right endpoint \(j\). If we compute the sum for each and check divisibility by \(K\), the overall time complexity is \(O(N^3)\), or \(O(N^2)\) even with prefix sums. Given the constraint \(N \leq 2 \times 10^5\), this is too slow to meet the time limit.
Prefix Sums and Remainder Properties
To efficiently handle contiguous subsequence sums, we use prefix sums. Let \(S_i\) be the total weight of the first \(i\) items. Then the sum from item \(i+1\) to item \(j\) can be expressed as \(S_j - S_i\).
The condition that this sum is a multiple of \(K\) can be written as: $\((S_j - S_i) \equiv 0 \pmod K\)\( Rearranging: \)\(S_j \equiv S_i \pmod K\)\( In other words, the problem can be rephrased as **"count the number of pairs \)(i, j)\( where the prefix sums have the same remainder when divided by \)K\("** (where \)0 \leq i < j \leq N$).
Efficient Counting
If there are \(c\) indices with the same prefix sum remainder, the number of ways to choose 2 from them is \({}_c\mathrm{C}_2 = \frac{c(c-1)}{2}\). By summing this value over all remainders (from \(0\) to \(K-1\)), we can compute the answer in \(O(N)\).
Algorithm
- Compute prefix sums: Scan the array from the beginning, computing the remainder of the current prefix sum divided by \(K\).
- Record the frequency of each remainder:
Use a hash map (in Python,
collections.Counteror a dictionary) to count how many times each remainder appears.- Note: You must include the initial prefix sum of 0 (the state where no items have been selected yet) in the count.
- Compute combinations: For each remainder with frequency \(c\), compute \(\frac{c(c-1)}{2}\) and sum up all these values.
Complexity
- Time complexity: \(O(N)\)
- Scanning the sequence takes \(O(N)\), and hash map operations take \(O(1)\) on average, so the overall complexity is \(O(N)\).
- Space complexity: \(O(N)\)
- In the worst case, all prefix sum remainders are distinct, so the hash map may store up to \(N+1\) elements.
Implementation Notes
Size of \(K\): \(K\) can be as large as \(10^9\), so attempting to count using an array would cause memory overflow. Be sure to use an associative array (dictionary).
Handling the initial value: Don’t forget to pre-count the prefix sum \(S_0 = 0\) (i.e.,
remainder_counts[0] = 1). If you forget this, subsequences starting from the beginning (e.g., when \(A_1 + A_2\) is a multiple of \(K\)) will not be counted correctly.Source Code
import sys
from collections import Counter
def solve():
# 入力を取得
input_data = sys.stdin.read().split()
if not input_data:
return
n = int(input_data[0])
k = int(input_data[1])
a = list(map(int, input_data[2:]))
# 累積和の余りをカウントする
# S_0 = 0 なので、初期値として 0 を 1 つカウントしておく
remainder_counts = Counter()
remainder_counts[0] = 1
current_sum = 0
for x in a:
current_sum = (current_sum + x) % k
remainder_counts[current_sum] += 1
# 同じ余りを持つ累積和のペアの数を計算する
# 余りが同じ S_i, S_j (i < j) があれば、その間の区間の和は K の倍数
ans = 0
for count in remainder_counts.values():
if count >= 2:
ans += count * (count - 1) // 2
print(ans)
if __name__ == "__main__":
solve()
This editorial was generated by gemini-3-flash-preview.
posted:
last update: