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

f(x) を「中央値をxとする部分文字列の数」とする。f(m)が解。

 g(y) = \sum_{y \leq x}f(x)とおくと

 f(m) = g(m) - g(m+1)を得る。

あとは、g(y)を求めるだけ。g(y)の値は、中央値がy以上の部分文字列の数を表す。 部分文字列の中央値がy以上である必要十分条件

(y以上の値の数)>=(y未満の値の数) 

とわかる。

よって

h(y, z) := 「部分文字列a[0, z)内での(y以上の値の数)-(y未満の値の数)」

とすれば、 h(y, r) - h(y, l) >= 0, l<rとなるような組(l, r)を数えればいいことがわかる。なので、BITなどを使えば簡単に求まる。

ll f(vi a, int b) {
    int n = sz(a);
    rep(i, n) {
        if (a[i] < b)a[i] = -1;
        else a[i] = 1;
    }
    ll re = 0;
    BIT bit(2*n+2);
    int s = 0;
    rep(i, n) {
        bit.add(n + s, 1);
        s += a[i];
        re += bit.sum(0, n + s);
    }
    return re;
}

int main() {
    int N, M;
    cin >> N >> M;
    vi A(N);
    rep(i, N)cin >> A[i];
    ll re = f(A, M) - f(A, M+1);
    cout << re << endl;
}