E2. Median on Segments (General Case Edition)

题意:E1的强化版本。问中位数是m的区间有多少个

思路:定义run(m): 中位数<=m的区间个数,则有式①:cnt[小于等于m的数] >= cnt[大于m的数] 。预处理一下,用数状数组维护。想到好简单。留个思考题吧,能不能用数状数组维护中位数>=m的个数,然后类似上面做呢? 结果是不太方便。因为长度奇偶的原因,当run(m)定义为: 中位数>=m的区间个数,式①的符号虽然改变了,但是等号取不取和长度奇偶有关系,这样问题就不统一了

#include<bits/stdc++.h>
#define PI acos(-1.0)
#define pb push_back
#define F first
#define S second
using namespace std;
typedef long long ll;
const int N=2e5+500;
const int MOD=1e9+7;
ll a[N],sum[N],tree[N*2+2];
ll n,m;
ll query(ll pos){
    ll res=0;
    for(ll i=pos;i>0;i-=i&(-i))   res+=tree[i];
    return res;
}

void update(ll pos){
    for(ll i=pos;i<=2*N;i+=i&(-i))    tree[i]+=1;
}
ll x[N],y[N];
ll run(ll m){ /// the number of subarray that median is <= m
    memset(tree,0,sizeof tree);
    memset(x,0,sizeof x);
    memset(y,0,sizeof y);
    memset(sum,0,sizeof sum);
    ll res=0;
    ll small=0,big=0;
    for(int i=1;i<=n;i++){
        if(a[i]>m) big++,x[i]=big,y[i]=y[i-1];
        else    small++,x[i]=x[i-1],y[i]=small;
    }
    for(int i=1;i<=n;i++)   sum[i]=y[i]-x[i];
    update(n+50);
    for(int r=1;r<=n;r++){
        res+=query(sum[r]+n+50);/// sum[r]-1 is a pos
        update(sum[r]+n+50);

    }
    return res;
}

int main(void){
    ios::sync_with_stdio(false);
    cin.tie(0);cout.tie(0);

    cin >>n>>m;
    for(int i=1;i<=n;i++)   cin >> a[i];
    cout << run(m)-run(m-1)<< endl;

    return 0;
}