PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: raysh07
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Dynamic programming
PROBLEM:
For a permutation , let f(P) denote the set of elements that appear as the median of some subarray of having length \gt 1
You’re given an integer N. Compute the number of distinct values f(P) can take across all permutations P of \{1, 2, \ldots, N\}.
EXPLANATION:
One property we’ll use a lot here is the following: if x is the median of some subarray of a permutation P, then there always exists a choice of subarray that either starts with or ends with x.
(To see why: read through the editorial to this problem first).
Now, we need to characterize subsets that can be the median-sets of some permutation.
To do that, it’s actually helpful to concentrate on the elements that are not medians.
Let’s consider an element x at position i, and suppose x is not a median of any subarray having length \gt 1.
Then, for each index j such that |i-j| \le 2, we must have P_j \lt x.
This is because:
- If an element adjacent to x is larger than it, taking x with that element results in x being the median.
- If P_{i+1} \lt x but P_{i+2} \gt x, the subarray [i, i+2] has median x, which is bad.
In particular, observe that this means any pair of non-medians must be at least three positions away from each other; or else the smaller one among them will have to be a median.
Using the above conclusion, let’s try to decide if some subset S of elements can all be made non-medians.
First off, note that N \in S for sure; since N definitely cannot be the median of any subarray when all the other elements are \lt N.
Otherwise, the observation about length 3 gives rise to the following greedy construction:
- Let x_1, x_2, \ldots, x_k be the elements of S, in ascending order.
- Let y_1, y_2, \ldots, y_{N-k} be the elements outside of S, in ascending order.
- Then, we have the following construction: place x_1, then y_1 and y_2, then x_2, then y_3 and y_4, then x_3, and so on.
Essentially, the idea is to separate the non-medians as little as possible. - The only exception to this is the element N: this alone should be placed at the very end of the array (to guarantee that y_{N-k} becomes a median).
Interestingly, it turns out that this greedy construction is optimal: in the sense that if it works then we’ve found a valid permutation, and if it doesn’t work then no permutation exists which can make S be the set of non-medians.
A proof of this will follow later, after we derive a couple more conditions - for now, we proceed assuming this is true.
Let’s analyze the construction a bit to see when it can be valid.
x_1 is placed first, and then y_1 and y_2.
So, we certainly need y_1 and y_2 to be both smaller than x_1.
Next comes x_2, which is \gt x_1 (and so larger than y_1 and y_2 as well).
After x_2 are y_3 and y_4, which again must both be smaller than x_2.
In general, observe that for each 1 \le i \lt k, x_i must be larger than y_{2i-1} and y_{2i}; which in turn also makes it larger than every element in the prefix before it too.
So, each non-median (other than N) must be able to be ‘matched’ with two medians smaller than it.
(We exclude i = k from this criterion because x_k = N is larger than every element anyway; and in particular doesn’t need two smaller medians after it to separate it from anything else).
It’s easy to see that this condition is also quite necessary, from the fact that all non-medians must be separated by a distance of at least 3.
Now, observe that every y_i will surely be a median (because each of them will always be adjacent to a larger element).
So, we only need to worry about checking if each x_i is a non-median or not.
Here, we use the fact that if some subarray median is x_i, then there will be a subarray either starting with or ending with x_i.
Since each x_i is strictly larger than everything to its left, it cannot be the median of any subarray ending at x_i.
Thus, we only look to the right of x_i.
Note that all elements larger than x_i lie to its right - there are N-x_i such elements.
Also, some elements smaller than x_i lie to the right of x_i.
Specifically, there are exactly 3\cdot (i-1) smaller elements to the left of x_i - this means there are x_i - 1 - 3\cdot (i-1) smaller elements to its right.
Now, observe that if N - x_i \ge x_i - 1 - 3\cdot (i-1), x_i will definitely be the median of some subarray starting with it (because, from the +1/-1 perspective, the sum of elements after x_i is non-negative so it’s possible to find a prefix with sum 0).
So, we definitely need N - x_i \lt x_i - 1 - 3\cdot (i-1) to hold.
We now know that for our construction to work, the i-th smallest non-median must satisfy the inequality N - x_i \lt x_i - 1 - 3\cdot (i-1).
It is not hard to prove that if this condition is satisfied, our construction will always create a valid permutation, since we’re placing elements in ascending order (again, looking at the +1/-1 perspective, it can be shown that if the sum ever goes non-negative there are too many elements larger than x_i, which contradicts the inequality).
Interestingly, this condition is not just sufficient - it’s also necessary!
Proof
We want to prove that in any subset of non-medians, the k-th largest median being x means N + 3k < 2x + 2 must hold.
Let’s fix N, and suppose some (x, k) pair violates the condition.
We’ll try to construct a permutation and see what happens.Note that we can safely assume that there are no non-medians \gt x (other than N); because if multiple k violate the condition for some subset, we can just choose the largest violating k.
Let’s place all the k non-medians (other than N) in some order, along with some smaller median elements between them to separate them.
Note that we’ll use at least (3k-3) elements that are \lt x for these k non-medians; so there are at most (x-1) - (3k-3) smaller elements remaining.Let the leftmost non-median be at index L, rightmost be at index R.
Observe that it’s impossible for us to have placed any elements \gt x between L and R, because of the following property: it can be shown that if i < j are positions of consecutive non-medians, then for all i < k < j we must have P_k \lt \max(P_i, P_j).
(This can be proved by contradiction by looking at the maximum element between P_i and P_j, and noting that if it becomes a median then either P_i or P_j will also become a median).Now, consider the prefix ending at L and the suffix starting at R.
In total, there are N-x elements \gt x that will appear in this prefix/suffix, and at most (x-1) - (3k-3) elements \lt x.
Since N-x \ge (x-1) - (3k-3), this means either the prefix or the suffix will have non-negative sum with respect to x from the +1/-1 perspective (and so with respect to all elements \lt x as well); which in turn means that either the element or L or the one at R will be a median, resulting in a contradiction.This takes care of the case where N is placed outside [L, R].
It’s possible that N can be placed in the range [L, R] too; and then we can’t invoke the property of “all elements \lt x must be outside [L, R]”, since some of these elements can be between N and some other non-median too.
However, with a bit of algebra, it can be seen that we run into trouble in this case as well; because inserting N in the middle requires using up two more numbers \lt x, and it can then be proved that to avoid any complications involving elements \gt x with the non-medians \lt x near N, we’ll have to use up more elements \lt x than elements \gt x.
In either case, either the middle will be bad, or one of the prefix/suffix will be bad, as before.
Now it’s trivial to write a dynamic programming solution that counts the number of valid subsets.
Let dp_{i, j} denote the number of ways to choose j non-medians from the elements \{1, 2, \ldots, i\}.
Transitions are simple:
- If i cannot be the j-th non-median, i.e. if it fails the check of
N - i \lt i - 1 - 3\cdot (j-1),
then dp_{i, j} = dp_{i-1, j}. - Otherwise, dp_{i, j} = dp_{i, j} + dp_{i-1, j-1}, since either we choose i to be the j-th non-median or we don’t.
The final answer is the sum of dp_{N-1, j} across all j, since the final non-median is forced to be N anyway.
TIME COMPLEXITY:
\mathcal{O}(N^2) per testcase.
CODE:
Editorialist's code (C++)
// #include <bits/allocator.h>
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());
int main()
{
ios::sync_with_stdio(false); cin.tie(0);
const int mod = 998244353;
int t; cin >> t;
while (t--) {
int n; cin >> n;
vector<int> dp(n+1);
dp[0] = 1;
for (int x = 1; x < n; ++x) {
for (int k = n; k >= 1; --k) {
if (n + 3*k < 2*x + 2) dp[k] = (dp[k] + dp[k-1]) % mod;
}
}
int ans = 0;
for (int i = 0; i <= n; ++i) {
ans = (ans + dp[i]) % mod;
}
cout << ans << '\n';
}
}