Solution Notes: If we replace each number less than X with -1 and each number greater than or equal to X with +1, then this problem reduces to counting the number of subarrays with nonnegative sums. There are a few ways to approach this. First, suppose we instead consider the complementary problem of counting the number of subarrays with negative sums. If we precompute an array P of prefix sums (so P[j] gives the sum of A[1...j]), subrrays A[i+1..j] with negative sums correspond to pairs (i,j) with i < j but P[i] > P[j]. Such pairs are called "inversions" in P, and counting inversions is a relatively standard algorithmic problem; this can be done in O(n log n) time, for example, via divide and conquer using a modified merge sort. In fact, with this particular problem, we can count the inversions in just O(n) time, as we will see shortly, owing to the fact that P[i+1] differs from P[i] by either -1 or +1.

A relatively simple O(n log n) solution due to Nathan Pinsker is shown below. It counts inversions using a binary index tree -- a useful data structure for many contest problems since it can be coded very quickly. The binary index tree implicitly encodes an array A[1..n] and supports two operations: query(p), which returns the sum of the prefix A[1..p], and update(p), which changes the value of A[p] (in our case here, we only need to increment A[p]).

``````
#include <cstdio>

using namespace std;
#define MAXN 100005

int n, x, a[MAXN], b[2 * MAXN], s;
long long total;

FILE *in = fopen("median.in", "r"), *out = fopen("median.out", "w");

int query(int p) {
int t = 0;
for (int i=(p + MAXN); i; i -= (i & -i)) {
t += b[i];
}
return t;
}
void update(int p) {
for (int i=(p + MAXN); i < 2*MAXN; i += (i & -i)) {
b[i]++;
}
}

int main() {
fscanf(in, "%d%d", &n, &x);
for (int i=0; i<n; ++i) {
fscanf(in, "%d", &a[i]);
a[i] = (a[i] >= x ? 1 : -1);
}
update(0);

for (int i=0; i<n; ++i) {
s += a[i];
total += query(s);
update(s);
}
fprintf(out, "%lld\n", total);

return 0;
}
``````
An even more concise O(n) solution is below, based on an approach suggested by Stan Zhang. Here, the idea is to count inversions as we go, as before, but to update our count somewhat carefully at each step. For each j, we want to count the number of indices i < j for which P[i] > P[j]. To do so, we keep a running count in b[x+MAXN] of the number of indices i < j with P[i] = x (that is, we keep a running count of how many times each prefix sum has occurred so far). Letting t be the number of indices i < j for which P[i] > P[j], we can update t as j moves to j+1 by either adding or subtracting an element of this b array, owing to the fact that P[j+1] is either one higher or one lower than P[j]. In the code below, s keeps track of our running prefix sum.
``````
#include <cstdio>
using namespace std;
#define MAXN 100005

int n, x, a, b[2 * MAXN], s = MAXN;
long long total, t;

FILE *in = fopen("median.in", "r"), *out = fopen("median.out", "w");

int main() {
fscanf(in, "%d %d", &n, &x);
total = (long long)n*(n+1)/2;
b[MAXN] = 1;
for (int j=0; j<n; ++j) {
fscanf(in, "%d", &a);
if (a >= x) { s++; t-=b[s]; }
else        { t+=b[s]; s--; }
b[s]++;
total -= t;
}
fprintf(out, "%lld\n", total);
return 0;
}

``````