CS Academy Subarray Medians

リンク : https://csacademy.com/contest/archive/task/subarray-medians/

問題概要

長さNの数列が与えられる。(1<=i<=j<=N かつ j-i=0 mod2)のすべてのi,jに関して[i,j]の中央値をmとしたとき、ijmの総和を求めよ。

解説

最初、iを決めて、前から順にBIT上のa[j]に1を足して、二部探索してO(N² logn)で求めたら、TLEしてしまった。ここで違うアプローチを考える。すべてのi,jの中央値を考えるのではなく、iがmとして取られるときの総和を求めることにする(割とこの考えは使う)。i<=jのとき、[i,j]に(a[i]より大きい個数)-(a[i]より小さい個数)をdifとする。d[dif]+=jをする。次にj<=iを見た時に、[j,i]に(a[i]より大きい個数)-(a[i]より小さい個数)をdifとすると、d[-dif]ja[i]をとることによってj番目から始まる区間で、中央値がi番目となるときをすべて考えることができる。

LL n, ans, a[11000], d[11000];
int main() {
    cin >> n;
    rep(i, n) {
        cin >> a[i];
    }
    rep(i, n) {
        rep(j, 11000)d[j] = 0;
        LL dif = 5000;
        for (int j = i; j >=0; j--) {
            if (a[i] > a[j])dif++;
            else dif--;
            if (dif <= 10000 && dif >= 0)d[dif] += (j + 1);
        }
        dif = 5000;
        for (int j = i; j < n; j++) {
            if (a[i] >= a[j])dif--;
            else dif++; 
            if (dif <= 10000 && dif >= 0)ans += a[i] * d[dif] * (j + 1);
        }
    }
    cout << ans << endl;
    return 0;
}