code/MeiTuan2/question4.cpp

/*求区间众数和
第一行输入一个正整数n。第二行输入n个正整数，1<n<200000,1<=a<=2
输入
3
2 1 2
输出
9

*/
#include <cstdio>
#include <iostream>
#include <algorithm>
typedef long long LL;
const int N = 200010;
int n;
int a[N], tr[N * 2];
int sum[N];
int lowbit(int x) {
    return x & -x;
}
void modify(int x, int k) {
    for (int i = x; i <= 2 * n + 5; i += lowbit(i)) tr[i] += k;
}
int query(int x) {
    int res = 0;
    for (int i = x; i; i -= lowbit(i)) res += tr[i];
    return res;
}
int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) scanf("%d", &a[i]);
    LL ans = 0;
    modify(0 + n + 1, 1);
    for (int i = 1; i <= n; i++) {
        sum[i] = sum[i - 1];
        if (a[i] == 1) sum[i] += 1;
        else sum[i] -= 1;
        ans += query(sum[i] + n + 1) + (i - query(sum[i] + n + 1)) * 2;
        modify(sum[i] + n + 1, 1);
    }
    printf("%lld", ans);
    return 0;
}