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

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

AtCoder ABC 146 E - Rem of Sum is Num (1D, 青色, 500 点)

AtCoder AtCoder500点 ABC-E 青色diff NoviSteps1D 式変形 f(i,j)をiとjとに分離する 考察：操作・条件・目的関数を言い換える Zero-Sum Ranges 解空間:O(N^2)個の区間 解空間:O(N^2)通りの選択肢 区間スライド系問題 区間

ちょっと工夫が必要な感じ

問題概要

長さ N 正整数列 A_{1}, A_{2}, \dots, A_{N} と、正の整数 K が与えられる。

数列の区間であって、総和を K で割った余りと、区間内に含まれる要素の個数とが等しいものの個数を求めよ。

制約

  • 1 \le N \le 2 \times 10^{5}
  • 1 \le K, A_{i} \le 10^{9}

考えたこと

結構ややこしい。落ち着いて条件を整理しよう！！

まず、0-indexed で考えることにして、数列 A の累積和を S_{0}, S_{1}, \dots, S_{N} としよう。このとき、問題は次のように再解釈できる。

長さ N+1 の数列 S_{0}, S_{1}, \dots, S_{N} が与えられる。

  • 0 \le i \lt j \le N
  • S_{j} - S_{i} ã‚’ K で割った余りが j - i に一致する

という条件を満たす (i, j) の組の個数を求めよ。

ここで、もし j - i \lt K であることが保証されるならば、話は早くて、

S_{j} - S_{i} \equiv j - i \pmod K
⇔ S_{i} - i \equiv S_{j} - j \pmod K

を満たすような (i, j) を求めればよくて、map を用いて自然に解ける。

そして、 j - i \ge K となることがあり得たとしても、少し考えてみれば、結局、次の条件を満たす (i, j) の組の個数を求めればよいことがわかる。

  • i \lt j
  • j - i \lt K
  • S_{i} - i \equiv S_{j} - j \pmod K

これを求めるためには、 j = 0, 1, \dots, N に対して、

  • j - K \lt i \gt \lt j
  • S_{i} - i \equiv S_{j} - j \pmod K

を満たす i の個数を求め、その値を足していけばよい。この操作は、やはり map を用いてできる。計算量は O(N \log N) と評価できる。

コード

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

// Sj - Si ≡ j - i mod K かつ j - i < K
// Si - i ≡ Sj - j mod K かつ j - i < K
int main() {
    long long N, K;
    cin >> N >> K;
    vector<long long> A(N), S(N+1, 0);
    for (int i = 0; i < N; i++) {
        cin >> A[i];
        S[i+1] = S[i] + A[i];
    }

    long long res = 0;
    map<long long, long long> nums;
    for (int j = 0; j <= N; j++) {
        long long rj = ((S[j] - j) % K + K) % K;
        res += nums[rj];

        // S[j] を追加
        nums[rj]++;

        // j - i = K となる i を削除
        int i = j - (K - 1);
        if (i >= 0) {   
            long long ri = ((S[i] - i) % K + K) % K;
            nums[ri]--;
            if (nums[ri] == 0) nums.erase(ri);
        }
    }
    cout << res << endl;
}

を満たす i の個数を