E1. Median on Segments (Permutations Edition) | Codeforces Round #496 (Div. 3)

数列なので中央値mはちょうど1つだけある。mの位置に注目すると、条件を満たす部分文字列は、mから左へ0以上、右へ0以上の長さだけ伸ばしたものに限られる。

部分文字列がmの左側へl, 右側へrだけ伸ばしたものであるとする。また、m以外の値はmとの大小関係だけわかればよい。

q[i] = 1 (p[i] > mのとき)
q[i] = -1(p[i] < mのとき)

とおく。mの位置をtとおくと  f(a, b) = \sum_{i=a}^{b-1}q_iとおくと  0 \leq f(t-l, t) + f(t+1, t+1+r) \leq 1であれば条件を満たす。 これは、部分文字列の左側と右側合わせて、

(mより大きい値の数) - (mより小さい値の数) = 0 または 1

ということ。

少し変形して  - f(t-l, t) \leq f(t+1, t+1+r) \leq 1 - f(t-l, t) と得る。 よって事前にf(t +1, t+1+r)をrの総当たりで固定して数えておく。あとはf(t-l, t)も総当たりして、各lについて条件を満たす f(t+1, t+1+r) を数えればいい。

map<int, int> L, R;
    
int N, M;
cin >> N >> M;
vi A(N);
int p = -1;
rep(i, N) {
    cin >> A[i];
    if (A[i] == M)p = i;
}

L[0] = 1;
R[0] = 1;

int x = 0;
for (int i = p - 1; i >= 0; --i) {
    if (A[i] > M)x++;
    else if (A[i] < M)x--;
    L[x] += 1;
}
x = 0;
for (int i = p + 1; i < N; ++i) {
    if (A[i] > M)x++;
    else if (A[i] < M)x--;
    R[x] += 1;
}

ll re = 0;
each(l, L) {
    // l + r = 0
    if (R.count(-l.first)) {
        re += (ll)l.second * R[-l.first];
    }
    // l + r = 1
    if (R.count(1 - l.first)) {
        re += (ll)l.second * R[1 - l.first];
    }
}
cout << re << endl;