ひらめの日常

プログラミングと読書と

ABC146-E Rem of Sum is Num (500)

問題はこちら

atcoder.jp

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

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

考え方

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

 r := j + 1, l := i とすると、以下のように問題を言い換えることができる。
「mod K において、 0 \leq l \lt r \leq N となる  l, r に対し、  d_r - d_l = r - l となるものの個数を求めよ。」

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

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

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

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

 K \leq r - l の時には、要素の数が  K 以上になるので、queueに  d _ {index} - {index} の値を入れておき、queueのサイズが  K 以上になったら、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] - はまやんはまやんはまやん