Official

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


\(S_i=A_1+\dots+A_i\) と定めます。\(S_0=0\) とします。

このとき、 \(A_L+\dots+A_R=S_R-S_{L-1}\) であることから、「\(A_L+\dots+A_R\)\(K\) の倍数」は「\(S_{R}\)\(S_{L-1}\)\(K\) で割った余りが等しい」と同値です。

以上より、 \(S_i\)\(K\) で割った余りを \(S'_i\) とすると、問題は「\(S'_L=S'_R\) を満たす、 \(0 \leq L < R \leq N\) を満たす整数の組 \((L,R)\) の個数は?」と読み替えられます。

この問題は各値の \(S'\) への登場回数を連想配列(辞書)を用いて求ることで解くことができます。

計算量は \(O(N\log N)\) です。

実装例 (C++)

#include<bits/stdc++.h>
using namespace std;

int main(){
  int n, k;
  cin >> n >> k;
  vector<int>a(n);
  for(int i=0; i<n; i++) cin >> a[i];

  vector<int>s(n+1);
  for(int i=0; i<n; i++){
    s[i+1] = (s[i] + a[i]) % k;
  }
  map<int,int>counter;
  for(int i=0; i<n+1; i++){
    counter[s[i]]++;
  }

  long long ans=0;
  for(auto[_, v]: counter){
    ans += (long long) v * (v-1) / 2;
  }
  cout << ans << endl;
}

実装例 (Python)

from collections import Counter
N, K = map(int, input().split())
A = list(map(int, input().split()))

S = [0]
for a in A:
  S.append((S[-1] + a) % K)
counter = Counter(S)

ans = 0
for v in counter.values():
  ans += v * (v-1) // 2
print(ans)

posted:
last update: