Atcoder E - Meaningful Mean（线段树+思维）

题目链接：http://arc075.contest.atcoder.jp/tasks/arc075_c

题意：问数组a有多少子区间平均值为k

题解：一开始考虑过dp，但是显然不可行，其实将每一个数都减去k就不用求平均值了，然后就是求满足前缀和sum[r]-sum[l-1]>=0,有多少组就行了。有点像逆序对。

用线段树即可。

#include <iostream>
#include <cstring>
#include <cstdio>
#include <algorithm>
using namespace std;
typedef long long ll;
const int M = 2e5 + 10;
struct TnT {
    int pos;
    ll val;
}nn[M];
struct node {
    int l , r;
    ll sum;
}T[M << 2];
bool cmp(TnT a , TnT b) {
    if(a.val == b.val) return a.pos < b.pos;
    return a.val < b.val;
}
void push_up(int i) {
    T[i].sum = T[i << 1].sum + T[(i << 1) | 1].sum;
}
void build(int i , int l , int r) {
    int mid = (l + r) >> 1;
    T[i].l = l , T[i].r = r , T[i].sum = 0;
    if(l == r) return ;
    build(i << 1 , l , mid);
    build((i << 1) | 1 , mid + 1 , r);
    push_up(i);
}
void update(int i , int pos) {
    int mid = (T[i].l + T[i].r) >> 1;
    if(T[i].l == pos && T[i].r == pos) {
        T[i].sum++;
        return ;
    }
    if(mid < pos) update((i << 1) | 1 , pos);
    else update(i << 1 , pos);
    push_up(i);
}
ll query(int i , int l , int r) {
    int mid = (T[i].l + T[i].r) >> 1;
    if(T[i].l == l && T[i].r == r) {
        return T[i].sum;
    }
    push_up(i);
    if(mid < l) return query((i << 1) | 1 , l , r);
    else if(mid >= r) return query(i << 1 , l , r);
    else {
        return query(i << 1 , l , mid) + query((i << 1) | 1 , mid + 1 , r);
    }
}
int main() {
    ll n , k , a;
    scanf("%lld%lld" ,  &n , &k);
    for(int i = 1 ; i <= n ; i++) {
        scanf("%lld" , &a);
        a -= k;
        nn[i].pos = i;
        nn[i].val = nn[i - 1].val + a;
    }
    sort(nn + 1 , nn + 1 + n , cmp);
    ll ans = 0;
    build(1 , 1 , (int)n);
    for(int i = 1 ; i <= n ; i++) {
        update(1 , nn[i].pos);
        if(nn[i].val >= 0) ans++;
        if(nn[i].pos - 1 == 0) continue;
        ans += query(1 , 1 , nn[i].pos - 1);
    }
    printf("%lld\n" , ans);
    return 0;
}