Codeforces Round #496 (Div. 3) E2 - Median on Segments (General Case Edition)

E2 - Median on Segments (General Case Edition)

题目大意:给你一个数组，求以m为中位数的区间个数。

思路：很巧秒的转换，我们把<= m 数记为1, >m的数 记为-1， 求其前缀,  我们将问题转变成求以<= m 的数作为中位数的区间个数，

答案就变为ans(m) - ans(m - 1)，我们可以用上面求得的前缀用bit就能求出答案。

#include<bits/stdc++.h>
#define LL long long
#define fi first
#define se second
#define mk make_pair
#define pii pair<int, int>

using namespace std;

const int N = 4e5 + 7;
const int M = 1e5 + 7;
const int inf = 0x3f3f3f3f;
const LL INF = 0x3f3f3f3f3f3f3f3f;
const int mod = 1e9 +7;

int n, m, a[N], b[N];
LL val[N];

void modify(int x, int v) {
    for(int i = x; i < N; i += i & -i)
        val[i] += v;
}

LL sum(int x) {
    LL ans = 0;
    for(int i = x; i; i -= i & -i)
        ans += val[i];
    return ans;
}

LL cal(int m) {
    memset(val, 0, sizeof(val));
    for(int i = 1; i <= n; i++) {
        b[i] = (a[i] <= m ? 1 : -1);
    }

    for(int i = 1; i <= n; i++) b[i] += b[i - 1];

    modify(n + 1, 1);

    LL ans = 0;
    for(int i = 1; i <= n; i++) {
        ans += sum(b[i] + n + 1);
        modify(b[i] + n + 1, 1);
    }
    return ans;
}

int main() {

    scanf("%d%d", &n, &m);
    for(int i = 1; i <= n; i++) {
        scanf("%d", &a[i]);
    }
    
    printf("%lld\n", cal(m) - cal(m - 1));
    return 0;
}


/*
3
3 2
*/