부분수열(VUDU) (JUNGOL 2956)

https://jungol.co.kr/problem/2956?cursor=MTMsMiw0

---

아이디어

구간 [i+1,j]의 수들의 평균이 P 이상이라고 해봅시다. sum[i]를 A[1]+A[2]+...+A[i]라고 정의하면, (sum[j]-sum[i])/(j-i)>=p라고 정의할 수 있습니다. 이를 정리하면 sum[j]-pj>=sum[i]-pi이고, 이것은 인덱스에 의존하므로, 모든 인덱스 i에 대하여 sum[i]-pi를 저장한 후에, i<j이고 sum[i]-pi<=sum[j]-pj인 (i,j)의 쌍의 개수를 세그먼트 트리를 이용해서 구해주면 문제를 해결할 수 있습니다.

---

소스코드

#include <iostream>
#include <vector>
#include <algorithm>
#define ll long long
using namespace std;

int tree[4040404]={0};
ll Sum(ll s,ll e,ll node,ll l,ll r)
{
    if (e<l || r<s) return 0;
    if (l<=s && e<=r) return tree[node];
    ll mid=(s+e)/2;
    return Sum(s,mid,node*2,l,r)+Sum(mid+1,e,node*2+1,l,r);
}
void update(ll s,ll e,ll node,ll idx)
{
    if (idx<s || e<idx) return;
    tree[node]++;
    if (s==e) return;
    ll mid=(s+e)/2;
    update(s,mid,node*2,idx);
    update(mid+1,e,node*2+1,idx);
}

ll solve(vector <ll> v) //v에서 i<j이면서 v[i]<=v[j]인 (i,j) 개수 세기
{
    ll N=v.size();
    vector <ll> v2;

    for (ll x:v) v2.push_back(x);

    sort (v2.begin(),v2.end());
    v2.erase(unique(v2.begin(),v2.end()),v2.end());

    for (ll i=0;i<N;i++)
        v[i]=lower_bound(v2.begin(),v2.end(),v[i])-v2.begin()+1;

    ll res=0;
    for (ll i=0;i<N;i++)
    {
        res+=Sum(1,N,1,1,v[i]);
        update(1,N,1,v[i]);
    }

    return res;
}

ll arr[1010101]={0};

int main()
{
    ios_base::sync_with_stdio(false); cin.tie(NULL);
    ll N,i,P,sum=0;
    vector <ll> v;

    cin>>N;
    for (i=0;i<N;i++) cin>>arr[i];
    cin>>P;

    v.push_back(0);
    for (i=0;i<N;i++)
    {
        sum+=arr[i];
        v.push_back(sum-P*(i+1));
    }

    cout<<solve(v);
}