AtCoder: ABC146-E Rem of Sum is Num (500)

問題はこちら

atcoder.jp

長さ の整数列 と正の整数 が与えられる。 の空でない連続する部分列であって、要素の和を で割った余りが要素の数と等しくなるものの数を求めよ。

以下は0-indexで話を進める。

考え方

要素の和の余りに関する問題なので、とりあえず累積和を mod K で取り、式に起こしてみる。mod K における、index i までの累積和を とすると、以下のようになる。

• 区間 ] の和:
• 区間 ] の要素の数:

とすると、以下のように問題を言い換えることができる。
「mod K において、 となる に対し、 となるものの個数を求めよ。」

しかし、 を全探索すると計算量が になり間に合わない。そこで、片方の端点 を固定して、その条件を満たす の個数を数えることを考える。

ここでポイントは、 以下の二つ。

• 要素の数が 以上の時は、要素の和が必ず 以上になり、題意を満たすことはない。つまり、 である。
• の時、mod K において以下の式変形が成り立つ。

なので、今までに見た index における の値の個数をmapなどで管理すれば、map[d[r] - r] で条件を満たす の個数を数えられる。

の時には、要素の数が 以上になるので、queueに の値を入れておき、queueのサイズが 以上になったら、map[queue.front()]--; queue.pop()して要素を取り除いてあげる。

解答

#include <bits/stdc++.h>

using namespace std;
using ll = long long;

#define rep(i, n) for(ll i = 0; i < (ll)(n); i++)

int main() {
    ll n, k;
    cin >> n >> k;
    vector<ll> a(n), d(n + 1);
    
    rep(i, n) cin >> a[i];
    rep(i, n) {
        d[i + 1] = d[i] + a[i];
        d[i + 1] %= k;
    }
    map<ll, ll> mp;
    ll ans = 0;
    queue<ll> que;

    rep(r, n + 1) {
        ll t = (d[r] - r) % k;
        if (t < 0) t += k;
        
        ans += mp[t];
        mp[t]++;
        
        que.push(t);
        if (que.size() >= k) {
            mp[que.front()]--;
            que.pop();
        }
    }
    cout << ans << '\n';
    return 0;
}

Submission #8703311 - AtCoder Beginner Contest 146

参考

とても分かり易かったです。

Rem of Sum is Num [AtCoder Beginner Contest 146 E] - はまやんはまやんはまやん