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とおくと とおくと であれば条件を満たす。 これは、部分文字列の左側と右側合わせて、
(mより大きい値の数) - (mより小さい値の数) = 0 または 1
ということ。
少し変形して と得る。 よって事前にf(t +1, t+1+r)をrの総当たりで固定して数えておく。あとはf(t-l, t)も総当たりして、各lについて条件を満たすを数えればいい。
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;