牛客练习赛91 D.监狱逃亡 (树状数组)

牛客练习赛91 D.监狱逃亡

因为总共就\(3\)行，先计算三行的前缀和，我们假设第一行在\(i\)处下来，第二行在\(j\)处下来，那么就有\(sum[1][i]+sum[2][j]-sum[2][i-1]+sum[3][n]-sum[3][j-1]\ge 0\)

移项得到\(sum[2][j]-sum[3][j-1] \ge sum[2][i-1]-sum[1][i]-sum[3][n]\)，因为\(j\ge i\),我们可以将右边的值离散化插入树桩数组中，每次在树桩数组中查询不大于左边项的个数贡献给答案。

#include <bits/stdc++.h>
using namespace std;
#define PII pair<int,int>
#define fi first
#define se second
#define pb push_back
#define ll long long
#define ull unsigned long long
#define PLL pair<ll,ll>
const int N=10000010;
const int mod=1e9+7;
const int INF=0x3f3f3f3f;

int n;
ll a[4][N];
ll sum[4][N];
vector<ll> all;
ll c[N];

ll getall(ll x){
    return lower_bound(all.begin(),all.end(),x)-all.begin();
}

ll lowbit(ll x){
    return x&(-x);
}

void update(ll i,ll k){
    while(i<=(int)all.size()){
        c[i]=(c[i]+k)%mod;
        i+=lowbit(i);
    }
}

ll query(ll i){
    ll res=0;
    while(i){
       res=(res+c[i])%mod; 
       i-=lowbit(i);
    }
    return res%mod;
}


int main(){
    scanf("%d",&n);
    for(int i=1;i<=3;++i){
        for(int j=1;j<=n;++j){
            scanf("%lld",&a[i][j]);
            sum[i][j]=sum[i][j-1]+a[i][j];
        }
    }
    for(int i=1;i<=n;++i){
        all.pb(sum[2][i]-sum[3][i-1]);
        all.pb(sum[2][i-1]-sum[1][i]-sum[3][n]);
    }
    sort(all.begin(),all.end());
    all.erase(unique(all.begin(),all.end()),all.end());
    ll ans=0;
    for(int i=1;i<=n;++i){
        ll p1=getall(sum[2][i]-sum[3][i-1]);
        ll p2=getall(sum[2][i-1]-sum[1][i]-sum[3][n]);
        update(p2+1,1);
        ans=(ans+query(p1+1))%mod;
    }
    printf("%lld\n",ans);
    return 0;
}