PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
Dynamic Programming, prefix sums
PROBLEM:
You’re given a partially filled permutation P.
Let f(P) denote the prefix maximum array of P.
Across all ways to fill in P to obtain a permutation, count the number of distinct values of f(P).
EXPLANATION:
We’ll try to fill in the missing elements of P from left to right.
Suppose we’ve filled in the first i elements of P.
We only care about prefix maximums, so let the current prefix maximum be m.
Note that this means that every element among the first i indices is \le m, and also m itself appears among these indices.
Then, when deciding element i+1,
- If P_{i+1} \ne -1, the value at this index is fixed.
There are a three possibilities based on its value:- P_{i+1} \lt m
In this case, the prefix maximum will remain m. - P_{i+1} \gt m
In this case, the prefix maximum will change to P_{i+1} instead. - P_{i+1} = m
This case is actually impossible because the elements of P must be distinct; and we already know m appears among the first i elements.
- P_{i+1} \lt m
- Alternately, it’s possible that P_{i+1} = -1, in which case we can freely decide its value.
Observe that again, there are only two cases - because the actual element we place doesn’t matter, only whether it’s smaller or larger than m (of course, once again placing m itself is not possible.)
Using these observations, we can formulate a solution using dynamic programming.
Define dp(i, m) to be the number of distinct prefix maximum arrays such that we’ve filled in the first i elements of the permutation, and the i-th prefix maximum is m.
There are \mathcal{O}(N^2) states.
For the transitions, we consider the case of P_i \ne -1 and P_i = -1 separately.
First, consider P_i \ne -1.
Then,
- If m \lt P_i, we have dp(i, m) = 0 because the i-th prefix maximum must certainly be at least P_i.
- If m \gt P_i, then P_i doesn’t affect the prefix maximum array in any way.
So we have dp(i, m) = dp(i-1, m). - If m = P_i, it means P_i is the new prefix maximum.
Here, the prefix maximum at index i-1 can be anything smaller than m.
So, we have dp(i, m) = dp(i-1, 1) + dp(i-1, 2) + \ldots + dp(i-1, m-1).
Next, consider the case of P_i = -1.
To compute dp(i, m), we have the following options:
- Clearly, we cannot place an element larger than m at this index (since the prefix maximum would then exceed m.)
So, either we place m and obtain a new prefix maximum, or we already had a prefix maximum of m and extend it by placing something smaller than m. - First, consider the case where we place m here.
The prefix maximum till index i-1 can be anything smaller than m, so once again we have
dp(i-1, 1) + dp(i-1, 2) + \ldots + dp(i-1, m-1)
ways.- Note that if m is already fixed to a different position by the input P, i.e. m is not ‘free’, we aren’t actually allowed to place m here; so in such a case the contribution of this case to dp(i, m) would be 0 instead.
- Next, we have the case where we place an element smaller than m.
Note that this means that the prefix max till index i-1 must already be m, so there are dp(i-1, m) possible ways.
However, we need to actually be able to place an element smaller than m at index i.
The actual element placed doesn’t matter (once the prefix max is \ge m, all elements \lt m are functionally equivalent).
So,- Let e_i denote the number of empty spaces among the first i indices of P.
- Let c_m denote the number of ‘free’ elements among \{1, 2, \ldots, m\}.
An element is called ‘free’ if it doesn’t appear in the initial P. - Then, it’s possible to place an element smaller than m at index i if and only if c_m \ge e_i.
- So, if c_m \ge e_i, we can add dp(i-1, m) to dp(i, m).
Otherwise, we do not.
The two arrays c and e used in the last case can be precomputed (either naively in \mathcal{O}(N^2) or with prefix sums in \mathcal{O}(N).)
The only slow part of the solution is computing the sum dp(i-1, 1) + \ldots + dp(i-1, m-1) which potentially needs to be done for every m; every other transition takes \mathcal{O}(1) time.
To optimize this, note that the value we’re looking for is exactly a prefix sum of the dp(i-1, \cdot) list.
So, simply storing these prefix sums will allow for every transition to be done in \mathcal{O}(1) time - giving a solution that works in \mathcal{O}(N^2) overall.
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 p(n+1, 0);
for (int i = 1; i <= n; ++i) cin >> p[i];
vector mark(n+1, 0); // mark[x] = 1 iff x is present in p
for (int x : p) mark[max(0, x)] = 1;
mark[0] = 1;
vector pref = mark; // pref[i] = number of free elements in [1, i]
for (int i = 0; i <= n; ++i) {
pref[i] = !mark[i];
if (i) pref[i] += pref[i-1];
}
vector dp(n+1, 0);
dp[0] = 1;
int pmax = 0, empty = 0;
for (int i = 1; i <= n; ++i) {
if (p[i] != -1) {
// j < p[i]: dp[i][j] = 0
// j = p[i]: dp[i][j] = sum(dp[i-1][x]) for x < p[i]
// j > p[i]: dp[i][j] = dp[i-1][j]
int sm = 0;
for (int j = 0; j < p[i]; ++j) {
sm = (sm + dp[j]) % mod;
dp[j] = 0;
}
dp[p[i]] = sm;
}
else {
++empty;
int sm = 0;
for (int j = 0; j <= n; ++j) {
int val = 0;
if (mark[j] == 0) { // Try placing j
val = sm;
}
// don't place j
if (j >= pmax and empty <= pref[j]) val = (val + dp[j]) % mod;
sm = (sm + dp[j]) % mod;
dp[j] = val;
}
}
pmax = max(pmax, p[i]);
}
cout << dp[n] << '\n';
}
}