[CF 1042D] Petya and Array

[CF 1042D] Petya and Array - 1800

문제 설명 :

배열 a[1], a[2].. a[n]이 주어져 있다. 여기서 임의의 [l, r] pair를 골랐을 때 $a[l]+a[l+1]+...a[r] < t$ 을 만족하는 pair가 총 몇 쌍인지 세야한다.

문제 풀이 :

$a[l]+a[l+1]+..a[r] < t$ 식을 [1,n]까지의 누적합을 저장하는 누적합배열 $psum[]$ 배열을 활용하여 $psum[r] - psum[l-1] < t$ 로 변형할 수 있다. 그럼 결국에는 $psum[r] - t < psum[l-1]$을 구하는 문제랑 똑같아진다. 변수가 l, r 총 2개이기 때문에 l 이나 r 을 고정시켜놓고 나머지 다른 값만 log만에 구하면 된다. 나는 r을 고정시켜놓고 이를 만족시키는 l의 개수를 구해볼 것이다. 즉 고정된 r에 대해서 $psum[r] - t$ 보다 큰 값을 가지는 $l (1 <= l <= r)$의 값을 구하면 된다.
특정 구간에서 특정 값 x보다 큰(혹은 작은)값의 개수를 구할 때는 offline query와 좌표압축을 사용하면 된다. 좌표압축으로 미리 값이 어느 위치에 들어갈지 배정한 다음, 1~n 원소를 하나씩 보면서 ans에다가 query값을 더해주고 update해주면 된다.

배울 점 :

이 문제에서 r이 아니라 l을 고정하려면 $psum[l-1] + t > psum[r]$ 로 식을 변형하고 이를 만족하는 r (l <= r <= n)을 구해야한다. 그런데 이렇게 하면 위에서 설명한 offline query로 전처리하고 update하는 방식을 사용할 수 없다. 따라서 이런 문제를 풀 때는 항상 left부터 update할 수 있는 방향으로 전처리를 해야한다.

소스코드 :

#include<stdio.h>
#include<algorithm>
#include<vector>
#include<map>
using namespace std;
typedef long long ll;

const int SIZE = 200009;
int tree[SIZE*8],arr[SIZE];
ll psum[SIZE];
int startIdx=1;

int count(int node,int cl, int cr,int ql,int qr){
    if(cl > qr || ql > cr) return 0;
    if(qr >= cr && cl >= ql) {
        return tree[node];
    }

    int mid = (cl+cr)/2;
    return count(node*2,cl,mid,ql,qr) + count(node*2+1,mid+1,cr,ql,qr);
}

void update(int node){
    node = node + startIdx - 1;
    tree[node]++;
    int par = node / 2;
    while(par >= 1){
        tree[par] = tree[par*2] + tree[par*2+1];
        par /= 2;
    }
    return;
}

int main(){
    int n;    scanf("%d",&n);
    ll t;   scanf("%lld",&t);

    vector<ll> v;
    v.push_back(0);
    for(int i=1 ; i<=n ; i++){
        scanf("%d",&arr[i]);
        psum[i] = psum[i-1] + arr[i];
        v.push_back(psum[i]);
    }
    for(; startIdx < v.size() ; startIdx *= 2);

    sort(v.begin(),v.end());
    v.erase(unique(v.begin(),v.end()),v.end());
    map<ll,int> compress;
    for(int i=0 ; i<v.size() ; i++){
        compress[v[i]] = i + 1;
    }
    update(compress[0]);

    ll ans = 0;
    for(int i=1 ; i<=n ; i++){
        ll val = psum[i] - t;
        int idx = upper_bound(v.begin(),v.end(),val) - v.begin() + 1;

        ans += count(1,1,startIdx,idx,startIdx);
        update(compress[psum[i]]);
    }

    printf("%lld\n",ans);
    return 0;
}