PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Prefix sums, segment trees/fenwick trees
PROBLEM:
You’re given an array A.
For each 1 \le L \le R \le N, define a new array f(L, R) as follows:
- Add +1 and -1 alternately to elements at indices L to R.
- All elements outside [L, R] remain unchanged.
Count the number of pairs (L, R) such that f(L, R) has a different median from the original array.
EXPLANATION:
It’s useful here to use an alternate characterization of the median: the median of an array of length N is equal to the smallest integer x such that strictly more than \frac{N}{2} elements of the array are \ge x.
Using this definition, let’s try to count the number of pairs (L, R) such that performing the alternating +1/-1 on A[L, R] causes the median to increase.
The other case, where the median decreases, can be handled similarly.
Let M denote the original median of A.
We want the median to increase - meaning it must be \ge M+1.
Observe that it’s enough to simply guarantee that strictly more than \frac{N}{2} elements are \ge M+1, because that’s enough for the median to be \ge M+1 already.
We’ll see what could possibly cause this.
First, note that each element changes by at most 1.
So, if A_i \ge M+2, then no matter what happens to it, it will remain \ge M+1 after the operation.
Similarly, if A_i \le M-1 then no matter what it’ll remain \lt M+1 after the operation.
This leaves us with A_i = M and A_i = M+1.
The only changes we care about here are:
- If A_i = M receives a +1, becoming M+1.
This increases the count of elements \ge M+1 by 1.
Observe that for a fixed (L, R), this can happen if and only if L \le i \le R and (i-L) is even. - If A_i = M+1 receives a -1, becoming M.
This decreases the count of elements \ge M+1 by 1.
Observe that for a fixed (L, R), this can happen if and only if L \le i \le R and (i-L) is odd.
We now use the above information.
Suppose we fix the left endpoint L of the subarray. We’ll try to count the number of valid R.
Let C denote the number of elements that are \ge M+1 in the initial array.
For each i \ge L, define the profit of index i as follows:
- If (i-L) is even and A_i = M, index i has a profit of +1.
- If (i-L) is odd and A_i = M+1, index i has a profit of -1.
- In all other cases, index i has a profit of 0.
Now, for any R \ge L, observe that the number of elements that will be \ge M+1 upon choosing (L, R) will equal exactly C plus the sum of profits of all indices from L to R.
We want this count to be \gt \frac{N}{2}, and so a right endpoint R is valid if and only if the sum of profits from L to R equals at least (\frac{N}{2}+1-C).
This counting exercise can be done using a segment tree/fenwick tree.
First, note that the profit of each index depends only on the parity of L; and not the exact value of L.
So, there are only two different profit arrays: one each for odd/even L.
Let’s fix L to be odd for now - even L can be handled similarly.
With L being odd, compute the profit at each index.
Let this profit be B_i at index i.
Then, for each odd L, we want to count the number of R \ge L such that
To deal with this range sum, we use prefix sums.
Define P_i = B_1 + B_2 + \ldots + B_i, with P_0 = 0.
Then the above expression becomes
Observe that upon fixing L, the right side is just a constant.
So, for a fixed odd L, we want to count R \ge L such that P_R is not smaller than some constant.
This can be done using a segment tree built on values as follows:
- Consider an array F such that F_x denotes the number of indices i having P_i = x.
- For each L = 1, 3, 5, \ldots in order,
- Compute the sum F_m + F_{m+1} + \ldots
Here, m = \frac{N}{2} + 1 - C + P_{L-1}. - Then, decrement the values of F_x for x = P_L and x = P_{L+1}, since these are no longer valid choices of R for future values of L.
- Compute the sum F_m + F_{m+1} + \ldots
- To support this quickly, we need to maintain F using a data structure that supports point updates and range sums; which a fenwick tree or segment tree can handle in \mathcal{O}(N\log N).
- Note that F_x may need to be indexed with negative values of x as well.
However, an easy way to get around that is to note that we only care about indices in [-N, N] since each index has a profit in [-1, 1], so we can add an offset of N and map this to the range [0, 2N] instead.
This way, we’re able to solve for all odd L, the number of choices of R \ge L that cause the median to increase.
Simply repeat this again for even L with the appropriate profit array; and then do something similar for the case where the median decreases (in which case you’ll care about M and M-1 instead.)
TIME COMPLEXITY:
\mathcal{O}(N\log N) per testcase.
CODE:
Tester's code (C++)
#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
template<typename T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;
typedef long long int ll;
typedef long double ld;
typedef pair<int,int> pii;
typedef pair<ll,ll> pll;
#define fastio ios_base::sync_with_stdio(false); cin.tie(NULL)
#define pb push_back
#define endl '\n'
#define sz(a) (int)a.size()
#define setbits(x) __builtin_popcountll(x)
#define ff first
#define ss second
#define conts continue
#define ceil2(x,y) ((x+y-1)/(y))
#define all(a) a.begin(), a.end()
#define rall(a) a.rbegin(), a.rend()
#define yes cout << "Yes" << endl
#define no cout << "No" << endl
#define rep(i,n) for(int i = 0; i < n; ++i)
#define rep1(i,n) for(int i = 1; i <= n; ++i)
#define rev(i,s,e) for(int i = s; i >= e; --i)
#define trav(i,a) for(auto &i : a)
template<typename T>
void amin(T &a, T b) {
a = min(a,b);
}
template<typename T>
void amax(T &a, T b) {
a = max(a,b);
}
#ifdef LOCAL
#include "debug.h"
#else
#define debug(...) 42
#endif
/*
*/
const int MOD = 1e9 + 7;
const int N = 1e5 + 5;
const int inf1 = int(1e9) + 5;
const ll inf2 = ll(1e18) + 5;
template<typename T>
struct fenwick {
int n;
vector<T> tr;
int LOG = 0;
fenwick() {
}
fenwick(int n_) {
n = n_;
tr = vector<T>(n + 1);
while((1<<LOG) <= n) LOG++;
}
int lsb(int x) {
return x & -x;
}
void pupd(int i, T v) {
for(; i <= n; i += lsb(i)){
tr[i] += v;
}
}
T sum(int i) {
T res = 0;
for(; i; i ^= lsb(i)){
res += tr[i];
}
return res;
}
T query(int l, int r) {
if (l > r) return 0;
T res = sum(r) - sum(l - 1);
return res;
}
int lower_bound(T s){
// first pos with sum >= s
if(sum(n) < s) return n+1;
int i = 0;
rev(bit,LOG-1,0){
int j = i+(1<<bit);
if(j > n) conts;
if(tr[j] < s){
s -= tr[j];
i = j;
}
}
return i+1;
}
int upper_bound(T s){
return lower_bound(s+1);
}
};
void solve(int test_case){
ll n; cin >> n;
vector<ll> a(n+5);
rep1(i,n) cin >> a[i];
vector<ll> cnt(n+5);
rep1(i,n) cnt[a[i]]++;
// find median
ll med = -1;
ll curr_cnt = 0;
rep1(x,n){
curr_cnt += cnt[x];
if(curr_cnt >= ceil2(n,2)){
med = x;
break;
}
}
auto go = [&](ll lp, ll need, vector<ll> b){
vector<ll> p(n+5);
p[0] = n+1;
rep1(i,n) p[i] = p[i-1]+b[i];
// set<ll> st(all(p));
// debug(sz(st));
fenwick<ll> fenw(2*n+5);
// p[i]-p[l-1] >= small_need
// p[l-1] <= p[i]-small_need
ll res = 0;
rep1(l,n){
if((l&1) == lp){
fenw.pupd(p[l-1],1);
}
ll val = p[l]-need;
if(val >= 1){
res += fenw.sum(val);
}
}
return res;
};
// new_med < med
ll small_count = 0;
rep1(i,n) small_count += (a[i] < med);
ll target = ceil2(n,2);
ll small_need = target-small_count;
ll small_ans = 0;
rep(lp,2){
vector<ll> b(n+5);
rep1(i,n){
if(a[i] == med-1){
// -1 if at odd ind
if((i-lp+1)&1){
b[i] = -1;
}
}
else if(a[i] == med){
// +1 if at even ind
if(!((i-lp+1)&1)){
b[i] = 1;
}
}
}
small_ans += go(lp,small_need,b);
}
// new_med > med
ll big_count = 0;
rep1(i,n) big_count += (a[i] > med);
ll big_need = n-target+1-big_count;
ll big_ans = 0;
rep(lp,2){
vector<ll> b(n+5);
rep1(i,n){
if(a[i] == med){
// +1 if at odd ind
if((i-lp+1)&1){
b[i] = 1;
}
}
else if(a[i] == med+1){
// -1 if at even ind
if(!((i-lp+1)&1)){
b[i] = -1;
}
}
}
big_ans += go(lp,big_need,b);
}
ll ans = small_ans + big_ans;
cout << ans << endl;
}
int main()
{
fastio;
int t = 1;
cin >> t;
rep1(i, t) {
solve(i);
}
cerr << "RUN SUCCESSFUL" << endl;
return 0;
}