裏紙

ほぼ競プロ、たまに日記

CF 877 F - Ann and Books

問題

Problem - F - Codeforces

問題概要

長さnの数列aと整数kが与えられる。クエリ(l,r):「閉区間 [ l,r] 内の連続部分列で、その和がkになるものの個数を答える」がq個与えられるのでそれぞれに答えよ。

  •  1 \le n \le 10^5
  •  -10^9 \le a_i \le 10^9
  •  -10^9 \le k \le 10^9
  •  1 \le q \le 10^5
  •  1 \le l \le r \le n

イデア

prefix sumを使えば、ある区間の和はO(1)で取ることが出来る。以下ではp_iと表す。

Mo's Algorithm(詳しくは他の人の詳しいを解説を見たほうが良いが、要素の更新がない/クエリ先読み可/区間の端を1だけ伸縮させた時の値の変化を簡単に得られる 時に有効)を使って区間の左端/右端を1伸ばす/縮める時に答えがどう変化するかを更新していきながら、該当するクエリと区間が一致した時に答えを記入していく、という流れになる。

具体的に何が分かると嬉しいかを考えてみる。いま、閉区間 [ l,r] に対する情報を持っている状況を考える。

右端を+1

右端を右に動かしたいとする。その時に、答えの増加分は(新たに増える区間は必ず右端がr+1なので)、 p_{r+1} - p_i = k (i \in [l-1, r ] )を満たすiの個数に一致する。

式変形すると、 p_i = p_{r+1} - kとなり、これを満たすiが何個あるかは、i \in [l-1, r ]に対してmapなどで頻度をカウントしておけばこの更新をO(logn)で行える事がわかる。

そして、答えの増分を計算した後に、mapにp_{r+1}の分を足して、完了となる。

左端を-1

左端を左に動かしたいとする。その時に、答えの増加分は(新たに増える区間は必ず左端がl-1なので)、 p_i - p_{(l-1)-1} = k (i \in [l-1, r ] )を満たすiの個数に一致し、式変形すると、 p_i = p_{(l-1)-1} + kとなる。

右端を-1

右端を左に動かしたいとする。その時に、答えの減少分は(消滅する区間は必ず右端がrなので)、 p_r - p_i = k (i \in [l-1, r-1 ] )を満たすiの個数に一致する。

答えの減少分を計算する前にmapからp_rの分を除き、 p_i = p_r - kを満たすiの個数を数えて答えから引いておけばよい。

左端を+1

左端を右に動かしたいとする。その時に、答えの減少分は(消滅する区間は必ず左端がlなので)、 p_i - p_{l-1} = k (i \in [l, r ] )を満たすiの個数に一致し、式変形すると、 p_i = p_{l-1} + kとなる。よって、答えの減少分を計算する前にmapからp_{l-1}の分を除き、減少分を計算すれば良い。

このように、似たような操作で区間を1伸ばしたり縮めたりがO(logn)で出来るので、全体の計算量がO((n+q)\sqrt{n} logn)となるが、残念ながらこの問題はこれを通させてくれない。だが、このlogは頑張って座標圧縮すると消すことが出来る。

今回、クエリ全体で現れる数は、p_x , p_x + k , p_x - k3n+3通りしかないので、事前処理としてこれらとindexを対応させておくことで、mapの部分はただの配列に置き換えることができ、O((n+q)\sqrt{n})とできる。

実装(C++)

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define all(x) (x).begin(),(x).end()
#define pb push_back
#define fi first
#define se second

using pi = pair<int,int>;
using P = pair<pi,int>;

const int B = 320;
const int Q = 100000;
const int N = 100002;

ll p[N];
// p[i][-k,0,k]
int f[N][3];

// ((R,L),qid)
vector<P> query[B];
ll ans[Q];

int L=1, R=0;
int ct[3*N]={};
ll now = 0;

void pr(){
    ++R;
    now += ct[f[R][0]];
    ++ct[f[R][1]];
}

void mr(){
    --ct[f[R][1]];
    now -= ct[f[R][0]];
    --R;
}

void pl(){
    --ct[f[L-1][1]];
    now -= ct[f[L-1][2]];
    ++L;
}

void ml(){
    --L;
    now += ct[f[L-1][2]];
    ++ct[f[L-1][1]];
}

int main(){
    int n;
    ll k;
    cin >>n >>k;

    vector<ll> a(n);
    rep(i,n)
    {
        cin >>a[i];
        if(a[i]==2) a[i] = -1;
    }
    rep(i,n){
        ll t;
        cin >>t;
        a[i] *= t;
    }

    // prefix sum
    p[0]=0;
    for(int i=1; i<=n; ++i) p[i] = p[i-1]+a[i-1];

    // 座標圧縮
    vector<ll> kk({-k,0,k});
    vector<ll> uniq_p;
    rep(i,n+1){
        for(ll j:kk) uniq_p.pb(p[i]+j);
    }
    sort(all(uniq_p));
    uniq_p.erase(unique(all(uniq_p)), uniq_p.end());

    rep(i,n+1){
        rep(j,3){
            ll val = p[i]+kk[j];
            f[i][j] = lower_bound(all(uniq_p),val) - uniq_p.begin();
        }
    }

    // クエリを左端をキーにして平方分割
    int q;
    cin >>q;
    rep(i,q){
        int l,r;
        cin >>l >>r;
        query[l/B].pb({{r,l},i});
    }

    ++ct[f[0][1]];
    rep(i,B){
        sort(all(query[i]));

        for(const auto &qq:query[i]){
            int l = qq.fi.se, r = qq.fi.fi, qid = qq.se;
            while(R<r) pr();
            while(r<R) mr();
            while(L<l) pl();
            while(l<L) ml();
            ans[qid] = now;
        }
    }

    rep(i,q) cout << ans[i] << endl;
    return 0;
}