けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 105 D - Candy Distribution (水色, 400 点)

AGC 023 A - Zero-Sum Ranges とほぼ同じ問題になってる。ただし、バケットを愚直に確保すると 109 サイズになってしまうので map とか連想配列が必要になる。

問題へのリンク

問題概要

 N 個の整数  A_1, A_2, \dots, A_N の連続する部分区間のうち、その総和が  M の倍数となっているものを数え上げよ。

制約

  •  1 \le N \le 10^{5}
  •  2 \le M \le 20^{9}

考えたこと

「配列連続する区間」は、「累積和をとった配列上の 2 点」と同一視することができるという超頻出視点なん。

高難易度問題でもこういういもす法的な視点が本質的な問題もあって、

あたりも気付きにくいけどこの視点が決定的役割を果たす。

今回の問題で言えば、S を A の累積和として


A[left, right) が M の倍数
⇔ S[right] - S[left] が M の倍数
⇔ S[right] と S[left] を M で割ったあまりが等しい


というのを利用する。結局、

  1. A の累積和配列 S を考える (スタートの「0」を忘れないように)
  2. S の各要素について M で割った余りで分類し、各余りについて何個あるのかを計上する (ただし M <= 109 なので、登場しうる余りが高々 N+1 個であることを利用して連想配列なりを用意する)
  3. 各余りについて個数を p としたとき、p(p-1)/2 を加算する

という感じで解くことができる。

#include <iostream>
#include <vector>
#include <map>
using namespace std;

int main() {
  int N, M; cin >> N >> M;
  vector<long long> A(N);
  for (int i = 0; i < N; ++i) cin >> A[i];
  map<long long, long long> nums;
  long long sum = 0;
  nums[sum]++;
  for (int i = 0; i < N; ++i) {
    sum += A[i];
    nums[sum % M]++;
  }
  long long res = 0;
  for (map<long long, long long>::iterator it = nums.begin(); it != nums.end(); ++it) {
    res += it->second * (it->second - 1) / 2;
  }
  cout << res << endl;
}