Official

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

Gemini 3.0 Flash

Overview

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

  1. Compute prefix sums: Scan the array from the beginning, computing the remainder of the current prefix sum divided by \(K\).
  2. Record the frequency of each remainder: Use a hash map (in Python, collections.Counter or 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.
  3. 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: