PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: raysh07
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Segment trees
PROBLEM:
You’re given a permutation P of the integers 1 to N.
For each prefix of P, compute the number of distinct medians across all subarrays of length \gt 1.
EXPLANATION:
Below, when we talk about subarrays, we only consider those with length \gt 1.
Consider some integer x.
Observe that if x appears as the median of some subarray in P[1, i], then it will also appear as the median of the same subarray in P[1, i+1].
So, for each element x, let’s define b_x to be the smallest integer such that x appears as the median of some subarray of P[1, b_x].
If x never appears as a subarray median, we say b_x = N+1 instead.
Assuming we’re able to compute all the b_x values, observe that the answer for the prefix of length i is simply the count of x such that b_x \le i.
The latter is trivial to find quickly once all the b_x values are known, so we focus on finding only the b_x values themselves.
For a fixed x, when can x appear as the median of some subarray of P?
There’s a somewhat standard way of answering this.
Let’s create a new array D of length N as follows:
- If P_i \lt x, set D_i = -1.
- If P_i \gt x, set D_i = 1.
- If P_i = x, set D_i = 0.
Now, observe that:
- If there’s an odd-length subarray with x as its median, the corresponding subarray in D must have sum equal to 0.
- If there’s an even-length subarray with x as its median, the corresponding subarray in D must have sum equal to 1.
So, we’re interested in subarrays containing the position of x that have sum equal to 0 or 1.
Importantly, observe that it’s enough to only care about subarray that either start with or end at x.
Why?
This follows from the fact that all the elements are +1 or -1, other than the position of x itself which contains 0.
First, note that if some element adjacent to x is larger than it, we immediately obtain a length-2 subarray having sum 1, so our claim is true and we’re done.
Otherwise, suppose x is at position i, and we’ve found subarray [L, R] with sum 0 or 1, where L \le i \le R.
If i = L or i = R we’re done. So, suppose L \lt i \lt R.
Observe that either the subarray [L, i] or the subarray [i, R] must have non-negative sum, since together they have non-negative sum.
Without loss of generality, say [L, i] has non-negative sum.We assumed that the elements adjacent to x were less than it.
This means D_{i-1} = -1.Now, consider what happens when we start at position i and move to position L, adding up the values of D along the way.
We start at i, with a value of 0.
Our first move is to index i-1, with value -1, so we’re currently negative.
Each time we move one step left, our sum changes by either +1 or -1.
When we eventually reach L, our sum is non-negative by assumption.Because the change was only +1 or -1, and we were initially negative but ended non-negative, we must have touched 0 at some point.
This gives us a subarray (with length greater than 1) ending with x having sum 0, as claimed.
So, if x is at index i, our aim is only to find subarrays starting at/ending at index i with sum 1 or 0.
This is trivial in linear time, but running it for each x results in an \mathcal{O}(N^2) algorithm which is too slow.
To optimize this, observe that if we change x to x+1, the array D doesn’t really change that much: in fact, only the values corresponding to the values x and x+1 change at all.
This allows to avoid having to rebuild the entire array D, and instead just maintain it online with point updates.
That is, we have the following process:
- Start with x = 1, and construct the array D.
- Solve the problem for the current value of x, i.e. check if there are subarrays with sum 0 or 1 that start/end at the appropriate position.
- Then, move x to x+1, updating the two positions in the array D in the process.
The only slow part here is the second step: checking for a subarray with sum 0 or 1 starting from/ending at a certain index i.
Let’s understand how to check whether there’s a subarray with sum = 0 or 1 starting at index i.
Observe that rather than exactly a sum of 0 or 1, we only really need to care about the maximum sum of some subarray starting at index i.
This again follows from the fact that the elements we’re dealing with are +1 and -1 only: as long as the maximum sum starting at this index is \ge 0, we know that some shorter subarray will have sum exactly equal to 0 or 1 (using the same argument as the proof above: either we’ll immediately see a 1 and finish; or we’ll go to -1 but end up \ge 0 so we must hit 0).
Now, we’re interested in finding the smallest prefix for which x is the median of some subarray.
It can be seen that:
- If there’s a subarray ending at x with median x, then every prefix containing the value x is going to have it as its median anyway.
- If there are no subarrays ending at x with median x, then we instead only need to care about the shortest subarray starting from x with median x.
This means, given the array D and a position i, we’re interested in the shortest subarray starting at i with sum either 0 or 1.
Let’s try to binary search for this.
Suppose we’re checking the subarray [i\ldots R].
Note that we only care about the maximum prefix sum starting from i in this subarray - because if this maximum prefix sum is \ge 0 then as seen above some shorter prefix will have sum 0, whereas if all prefix sums are \lt 0 we need to look further than R.
For a fixed i and R, finding the maximum prefix sum starting at i in the range [i\ldots R] can be done in logarithmic time using a segment tree (store in each node the sum of the entire range, as well as the maximum prefix sum; merging is easy to figure out).
Along with binary search, this gives us a check in \mathcal{O}(\log^2 N) time, which can be optimized to \mathcal{O}(\log N) using segment tree descent.
Since moving from x to x+1 corresponds to only a couple of point updates on the array D, the segment tree can handle this too in logarithmic time.
We thus have a solution in \mathcal{O}(N\log N) or \mathcal{O}(N\log^2 N) overall, which is more than fast enough.
TIME COMPLEXITY:
\mathcal{O}(N\log N) or \mathcal{O}(N\log^2 N) per testcase.
CODE:
Tester's code (C++)
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF (int)1e18
mt19937_64 RNG(chrono::steady_clock::now().time_since_epoch().count());
struct segtree{
struct node{
int x = 0;
int pref = 0, suf = 0;
void apply(int l, int r, int y){
x = y;
pref = y;
suf = y;
}
};
int n;
vector <node> seg;
node unite(node a, node b){
node res;
res.x = a.x + b.x;
res.pref = max(a.pref, a.x + b.pref);
res.suf = max(b.suf, b.x + a.suf);
return res;
}
void push(int l, int r, int pos){
}
void pull(int pos){
seg[pos] = unite(seg[pos * 2], seg[pos * 2 + 1]);
}
void build(int l, int r, int pos){
if (l == r){
return;
}
int mid = (l + r) / 2;
build(l, mid, pos * 2);
build(mid + 1, r, pos * 2 + 1);
pull(pos);
}
template<typename M>
void build(int l, int r, int pos, vector<M> &v){
if (l == r){
seg[pos].apply(l, r, v[l]);
return;
}
int mid = (l + r) / 2;
build(l, mid, pos * 2, v);
build(mid + 1, r, pos * 2 + 1, v);
pull(pos);
}
node query(int l, int r, int pos, int ql, int qr){
push(l, r, pos);
if (l >= ql && r <= qr){
return seg[pos];
}
int mid = (l + r) / 2;
node res{};
if (qr <= mid) res = query(l, mid, pos * 2, ql, qr);
else if (ql > mid) res = query(mid + 1, r, pos * 2 + 1, ql, qr);
else res = unite(query(l, mid, pos * 2, ql, qr), query(mid + 1, r, pos * 2 + 1, ql, qr));
pull(pos);
return res;
}
template <typename... M>
void modify(int l, int r, int pos, int ql, int qr, M&... v){
push(l, r, pos);
if (l >= ql && r <= qr){
seg[pos].apply(l, r, v...);
return;
}
int mid = (l + r) / 2;
if (ql <= mid) modify(l, mid, pos * 2, ql, qr, v...);
if (qr > mid) modify(mid + 1, r, pos * 2 + 1, ql, qr, v...);
pull(pos);
}
segtree (int _n){
n = _n;
seg.resize(4 * n + 1);
build(1, n, 1);
}
template <typename M>
segtree (int _n, vector<M> &v){
n = _n;
seg.resize(4 * n + 1);
if (v.size() == n){
v.insert(v.begin(), M());
}
build(1, n, 1, v);
}
node query(int l, int r){
return query(1, n, 1, l, r);
}
node query(int x){
return query(1, n, 1, x, x);
}
template <typename... M>
void modify(int ql, int qr, M&...v){
modify(1, n, 1, ql, qr, v...);
}
};
void Solve()
{
int n; cin >> n;
vector <int> p(n + 1), pos(n + 1);
for (int i = 1; i <= n; i++){
cin >> p[i];
pos[p[i]] = i;
}
segtree seg(n);
int v1 = -1, v2 = +1;
for (int i = 1; i <= n; i++){
seg.modify(i, i, v2);
}
vector <int> ans(n + 1);
int ct = 0;
for (int i = 1; i <= n; i++){
int q = pos[i];
if (q > 1 && p[q] < p[q - 1]){
ans[q] += 1;
} else if (q > 1 && seg.query(1, q - 1).suf >= 0){
ans[q] += 1;
} else if (q < n && p[q] < p[q + 1]){
ans[q + 1] += 1;
} else if (q < n && seg.query(q + 1, n).pref >= 0){
++ct;
int lo = q + 1, hi = n;
while (lo != hi){
int mid = (lo + hi) / 2;
if (seg.query(q + 1, mid).pref >= 0){
hi = mid;
} else {
lo = mid + 1;
}
}
ans[lo] += 1;
}
seg.modify(q, q, v1);
}
for (int i = 1; i <= n; i++){
ans[i] += ans[i - 1];
cout << ans[i] << " \n"[i == n];
}
}
int32_t main()
{
auto begin = std::chrono::high_resolution_clock::now();
ios_base::sync_with_stdio(0);
cin.tie(0);
int t = 1;
// freopen("in", "r", stdin);
// freopen("out", "w", stdout);
cin >> t;
for(int i = 1; i <= t; i++)
{
//cout << "Case #" << i << ": ";
Solve();
}
auto end = std::chrono::high_resolution_clock::now();
auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
cerr << "Time measured: " << elapsed.count() * 1e-9 << " seconds.\n";
return 0;
}