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

問題へのリンク

問題概要

長さ $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$ の個数を

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

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

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

問題概要

制約

考えたこと

コード