PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: nikhil_010309
Tester: raysh07
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
DP, combinatorics
PROBLEM:
For an array X with elements in [1, M] and a permutation P of [1, M], define P(X) to be the array [P_{X_1}, P_{X_2}, \ldots, P_{X_N}].
Define P^k(X) to be this iterated k times.
Given the array A, and values K, D; count pairs (X, P) such that:
- X is an array of length N with values in [1, M]
- X \le A lexicographically
- P’s graph consists of K connected components.
- X \lt P(X) \lt P^2(X) \lt\ldots\lt P^D(X) lexicographically.
EXPLANATION:
From the easy version, which could ignore the X \le A condition, the solution was:
- X \lt P(X) \lt P^2(X) \lt\ldots\lt P^D(X) lexicographically depends only on the first index of X that’s not fixed by P.
- This allowed us to fix this first non-fixed index, and then do some not-so-hard combinatorics to do the counting.
All those observations still apply here, we just need to deal with X \le A.
As before, we fix the first index r such that P_{X_r} \ne X_r.
Now we need to deal with X \le A.
For X \le A to hold, X must fall below A at some point (or remain equal to it always.)
Given that r is already fixed, there are three possibilities: this fall can happen before r, at r, or after r.
We’ll deal with each case in turn.
Case 1: The fall happens before r.
Here, let’s define dp(r, k) to be the number of ways to choose the first r-1 elements such that the resulting prefix is smaller than the corresponding prefix of A, and also exactly k distinct elements are used.
This can be calculated in similar fashion to the Stirling dp: separate cases depending on whether the fall is at r-1 or before r-1, where the second case needs you to look at dp(r-1, k) and dp(r-1, k-1) while the first needs to know the number of distinct elements till r-2 only; and the choice of what to place at r-1 can be computed from that.
Once we’re lex smaller already though, we’re functionally back in the easy version - so with both r and k fixed the remaining computation is exactly the same as it was there.
This is handled in \mathcal{O}(NM).
Case 2: The fall happens at r.
Here, the prefix till r-1 is fixed, and so is the number of distinct elements there.
We can try all possible choices of X_r \lt A_r, and for each one do a similar computation as the easy version since there are no more restrictions.
One thing to be careful about here is choosing the chain of elements from X_r, since that depends on the number of free values available that are \gt X_r; this is easily handled by iterating over values of X_r in descending order.
The complexity is still \mathcal{O}(NM).
Case 3: The fall happens after r.
For this, note that X_r is itself fixed to be A_r (and this state is only valid if A_r hasn’t appeared at a smaller index in A before this.)
Here, again we can just choose some D elements larger than A_r (that don’t appear in the prefix before r), and then complete the permutation appropriately (using the Stirling count.)
However, we also need to complete A - and unlike the previous two cases, we can’t simply choose elements willy-nilly.
Instead, note that we need to choose a suffix of X that’s smaller than the corresponding suffix of A.
This count is easy to precompute on each suffix; since either we go smaller at the start of the suffix, or we stay equal there and then the remaining suffix must be smaller.
Thus, with precomputation this case can also be handled in \mathcal{O}(NM), completing the problem.
TIME COMPLEXITY:
\mathcal{O}(NM) 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());
const int mod = 998244353;
struct mint{
int x;
mint (){ x = 0;}
mint (int32_t xx){ x = xx % mod; if (x < 0) x += mod;}
mint (long long xx){ x = xx % mod; if (x < 0) x += mod;}
int val(){
return x;
}
mint &operator++(){
x++;
if (x == mod) x = 0;
return *this;
}
mint &operator--(){
if (x == 0) x = mod;
x--;
return *this;
}
mint operator++(int32_t){
mint result = *this;
++*this;
return result;
}
mint operator--(int32_t){
mint result = *this;
--*this;
return result;
}
mint& operator+=(const mint &b){
x += b.x;
if (x >= mod) x -= mod;
return *this;
}
mint& operator-=(const mint &b){
x -= b.x;
if (x < 0) x += mod;
return *this;
}
mint& operator*=(const mint &b){
long long z = x;
z *= b.x;
z %= mod;
x = (int)z;
return *this;
}
mint operator+() const {
return *this;
}
mint operator-() const {
return mint() - *this;
}
mint operator/=(const mint &b){
return *this = *this * b.inv();
}
mint power(long long n) const {
mint ok = *this, r = 1;
while (n){
if (n & 1){
r *= ok;
}
ok *= ok;
n >>= 1;
}
return r;
}
mint inv() const {
return power(mod - 2);
}
friend mint operator+(const mint& a, const mint& b){ return mint(a) += b;}
friend mint operator-(const mint& a, const mint& b){ return mint(a) -= b;}
friend mint operator*(const mint& a, const mint& b){ return mint(a) *= b;}
friend mint operator/(const mint& a, const mint& b){ return mint(a) /= b;}
friend bool operator==(const mint& a, const mint& b){ return a.x == b.x;}
friend bool operator!=(const mint& a, const mint& b){ return a.x != b.x;}
mint power(mint a, long long n){
return a.power(n);
}
friend ostream &operator<<(ostream &os, const mint &m) {
os << m.x;
return os;
}
explicit operator bool() const {
return x != 0;
}
};
struct factorials{
int n;
vector <mint> ff, iff;
factorials(int nn){
n = nn;
ff.resize(n + 1);
iff.resize(n + 1);
ff[0] = 1;
for (int i = 1; i <= n; i++){
ff[i] = ff[i - 1] * i;
}
iff[n] = ff[n].inv();
for (int i = n - 1; i >= 0; i--){
iff[i] = iff[i + 1] * (i + 1);
}
}
mint C(int n, int r){
if (n == r) return mint(1);
if (n < 0 || r < 0 || r > n) return mint(0);
return ff[n] * iff[r] * iff[n - r];
}
mint P(int n, int r){
if (n < 0 || r < 0 || r > n) return mint(0);
return ff[n] * iff[n - r];
}
mint solutions(int n, int r){
// Solutions to x1 + x2 + ... + xn = r, xi >= 0
return C(n + r - 1, n - 1);
}
mint catalan(int n){
return ff[2 * n] * iff[n] * iff[n + 1];
}
};
const int PRECOMP = 3e6 + 69;
factorials F(PRECOMP);
// REMEMBER To check MOD and PRECOMP
void Solve()
{
int n, m, k, d; cin >> n >> m >> k >> d;
vector <int> a(n + 1);
for (int i = 1; i <= n; i++) cin >> a[i];
vector<vector<mint>> dp(m + 1, vector<mint>(m + 1, 0));
dp[0][0] = 1;
for (int i = 2; i <= m; i++){
for (int j = 1; j <= i / 2; j++){
dp[i][j] = (i - 1) * (dp[i - 1][j] + dp[i - 2][j - 1]);
}
}
int S = max(n, m);
// compute i^j, i^(-j) for all i, j <= S
vector<vector<mint>> pp(S + 1, vector<mint>(S + 1));
auto ipp = pp;
for (int i = 0; i <= S; i++){
pp[i][0] = 1;
ipp[i][0] = 1;
mint inv = 1 / mint(i);
for (int j = 1; j <= S; j++){
pp[i][j] = pp[i][j - 1] * i;
ipp[i][j] = ipp[i][j - 1] * inv;
}
}
mint ans = 0;
auto get = [&](int n, int k){
// n things, exactly one guy may be fixed, others all cycle >= 2 size
mint ans = 0;
if (n >= 1 && k >= 1 && k <= n){
ans += dp[n][k];
ans += dp[n - 1][k - 1];
}
return ans;
};
// case 1: x[p] is non-fixed, x[1, p] = a[1, p], x[p + 1, n] <= a[p + 1, n]
// x[1, p - 1] must be fixed
{
vector <bool> st(m + 1, false);
int distinct = 0;
for (int p = 1; p <= n; p++){
if (!st[a[p]]){
int val = a[p];
int options = 0;
for (int i = val + 1; i <= m; i++){
if (!st[i]){
options++;
}
}
mint ways = F.C(options, d);
mint suffix = 1;
for (int i = p + 1; i <= n; i++){
suffix += (a[i] - 1) * pp[m][n - i];
}
for (int fixed = 0; fixed < k; fixed++){
if (fixed < distinct) continue;
mint here = ways;
here *= F.C(m - d - 1 - distinct, fixed - distinct);
here *= get(m - fixed - d, k - fixed);
here *= suffix;
ans += here;
}
}
distinct += !st[a[p]];
st[a[p]] = true;
}
}
// case 2 : x[p] is non-fixed, x[p] < a[p], x[1, p - 1] = a[1, p - 1]
{
vector <bool> st(m + 1, false);
int distinct = 0;
for (int p = 1; p <= n; p++){
int opt1 = 0;
int opt2 = 0;
for (int i = 1; i <= m; i++) if (!st[i]){
if (i < a[p]) opt1++;
else opt2++;
}
// choose d + 1 elements such that the smallest is < a[p]
mint ways = F.C(opt1 + opt2, d + 1) - F.C(opt2, d + 1);
mint suffix = pp[m][n - p];
for (int fixed = 0; fixed < k; fixed++){
if (fixed < distinct) continue;
mint here = ways;
here *= F.C(m - d - 1 - distinct, fixed - distinct);
here *= get(m - fixed - d, k - fixed);
here *= suffix;
ans += here;
}
distinct += !st[a[p]];
st[a[p]] = true;
}
}
// case 3 : x[p] is unfixed, but x[q] < a[q] for some q < p
{
vector <bool> st(m + 1, false);
int distinct = 0;
vector <mint> store(k + 1, 0);
for (int fixed = 1; fixed < k; fixed++){
for (int p = 1; p <= n; p++){
store[fixed] += pp[fixed][p - 1] * pp[m][n - p];
}
}
for (int q = 1; q < n; q++){
// either x[q] exists in prefix, or x[q] is new
mint opt1 = 0, opt2 = 0;
for (int i = 1; i < a[q]; i++){
if (st[i]) opt1++;
else opt2++;
}
for (int fixed = 1; fixed < k; fixed++){
if (fixed < distinct) continue;
mint ways = 0;
{
mint here = F.C(m - distinct, fixed - distinct);
here *= F.C(m - fixed, d + 1);
here *= get(m - fixed - d, k - fixed);
ways += here * opt1;
}
if (fixed > distinct){
mint here = F.C(m - distinct - 1, fixed - distinct - 1);
here *= F.C(m - fixed, d + 1);
here *= get(m - fixed - d, k - fixed);
ways += here * opt2;
}
// compute fixed^(p - 1) * m^(n - p) for all p, fixed, store indexed by fixed
// then q = 1, ..., n removes a possible value of p + multiplies everything by fixed^(-q)
store[fixed] -= pp[fixed][q - 1] * pp[m][n - q];
mint mul = store[fixed] * ipp[fixed][q];
ans += ways * mul;
}
distinct += !st[a[q]];
st[a[q]] = true;
}
}
cout << ans << "\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;
}