円環上の Zero-Sum Ranges!!

問題へのリンク

問題概要

円周上に $N$ 個の地点 $1, 2, \dots, N$ がこの順に時計回りに並んでいる。地点 $i$ から地点 $i+1$ （ $i = N$ のとき $i + 1 = 1$ とする）まで時計回りに移動するのに要する時間は $A_{i}$ である。

次の条件を満たす $s, t$ の個数を求めよ

$1 \le s, t \le N$
$s \neq t$
$s$ から時計回りに $t$ へと到達するのに要する時間は $M$ の倍数である

制約

$2 \le N \le 2 \times 10^{5}$
$1 \le A_{i} \le 10^{9}$
$1 \le M \le 10^{6}$

考えたこと

まず、円周上の点を扱う一般的なテクニックとして、2 週するというのがある。今回は地点 $0, 1, \dots, N-1$ （0-indexed にする）があるわけだが、2 週させて地点 $0, 1, \dots, N-1, N, N+1, \dots, 2N-1$ を考えるのだ。ここで、地点 $N, N+1, \dots, 2N-1$ は地点 $0, 1, \dots, N-1$ に一致する。

そうすると、問題は次のように言い換えられる。

次の条件を満たす $s, t$ の個数を求めよ

$0 \le s \lt t \le 2N-1$
$t - s \lt N$
$s$ から時計回りに $t$ へと到達するのに要する時間は $M$ の倍数である

累積和を考える

さて、 $s$ から時計回りに $t$ へと到達するのに要する時間とは、つまり数列 $A$ の区間 $\lbrack s, t)$ の総和である。このように、数列のある区間の総和を考える際には累積和をとるとうまくいくことが多い。

S[v]：地点 0 から地点 $v$ まで時計回りに到達するのに要する時間（ただし $v \ge N$ のときは 1 週分の所要時間を足す）

こうすると、上記の 3 つめの条件は次のように言い換えられる。

S[s] ≡ S[t] (mod. $M$ )

こうしてみると、次の解法が考えられるだろう。

S[0], S[1], ..., S[2N-1] を $M$ で割った余りで分類する。そして、各 $r = 0, 1, \dots, M-1$ に対して、S[x] を $M$ で割った余りが $r$ であるような $x$ を集めた集合 $T_{r}$ を考える。

この集合 $T_{r}$ から 2 要素を選ぶ方法であって、その差が $N$ 未満であるものの個数を求める。

この個数を $r = 0, 1, \dots, M-1$ について足した値が答えである

なお、上記の個数を求める方法については、しゃくとり法でもいいし二分探索法でもよい。

各 $r$ に対する集合 $T_{r}$ に対して、しゃくとり法を用いた場合の計算量は $O(|T_{r}|)$ であり、 $\sum_{r = 0}^{M-1} |T_{r}| = O(N)$ であることから、全体の計算量は $O(N + M)$ となる。

コード

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

int main() {
    long long N, M;
    cin >> N >> M;
    vector<long long> A(N), S(N * 2 + 1, 0);
    for (int i = 0; i < N; i++) cin >> A[i];

    vector<vector<int>> where(M);
    for (int i = 0; i < N * 2; i++) {
        S[i + 1] = S[i] + A[i % N];
        where[S[i] % M].push_back(i);
    }

    long long res = 0;
    for (int r = 0; r < M; ++r) {
        const auto &v = where[r];
        int left = 0, right = 0;
        for (; left < v.size() && v[left] < N; left++) {
            while (right < v.size() && v[right] - v[left] < N) right++;
            res += right - left - 1;
        }
    }
    cout << res << endl;
}