PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07
Tester: iceknight1093
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
Inclusion-exclusion, combinatorics
PROBLEM:
Given K, a sequence of integers is called good if it contains some K consecutive integers.
For a permutation P, define f(P) to be the length of the shortest good prefix of P.
Given N and K, compute the sum of f(P) across all permutations of [1, N].
Here, N \le 2\cdot 10^5.
EXPLANATION:
To recap the solution from the easy version:
- Define c_x to be the number of subsets of \{1, 2, \ldots, N\} having size x that are not good.
- Then, the answer is\sum_{x=0}^N c_x \cdot x! \cdot (N-x)!
In the easy version, we were able to compute all the c_x values in \mathcal{O}(N^2) using dynamic programming.
However, that won’t quite work here, so we need a different approach.
Let’s try to solve for a fixed value of x using combinatorics.
We want to count the number of ways to pick x elements from \{1, 2, \ldots, N\} such that K consecutive elements are not chosen.
One way to look at this is that we have several disjoint segments representing the chosen elements; and each segment must have length \lt K.
Further, the total length of the segments must be x.
There are N-x elements that are not part of any segment.
Let’s call these “outside” elements.
Suppose we place all outside elements in a line.
Now, observe that each segment must be placed between two adjacent outside elements - or before the first/after the last outside element.
Further, we also can’t place multiple segments in the same location.
Thus, there are N-x+1 possible locations where the segments can be placed.
If we also allow for empty segments, then there are also exactly N-x+1 segments that will be placed too.
These N-x+1 segments must have a total length of x, but we also want each of their lengths to be \lt K.
To count the number of valid configurations, we use stars and bars along with inclusion-exclusion.
Let M = N-x+1 denote the number of segments.
Each segment must have non-negative length - let y_i denote the length of the i-th segment.
Our goal is to count tuples (y_1, y_2, \ldots, y_M) such that \sum y_i = x and 0 \le y_i \lt K for each i.
If we didn’t have the y_i \lt K restriction, then this would be simple: the number of ways to write x as the sum of M non-negative integers is given by stars and bars to be
Note that M+x=N+1 so this reduces to \binom{N}{M-1}.
Dealing with the y_i \lt K restriction is where we need inclusion-exclusion.
Let’s fix an integer j \le M and try to count the number of configurations in which j segments have length \ge K.
There are \binom{M}{j} ways to choose which j segments violate the condition.
Then, we can subtract K from the lengths of each of these j segments; after which there are no more restrictions: each segment must just have non-negative length.
The number of configurations is, again, given by stars-and-bars to be
(this is after applying the reduction of M+x=N+1.)
The subtraction of j\cdot K is because we removed K each from j segments, so the target sum reduces by the same amount.
Now, this is technically inaccurate because we didn’t place any restrictions on the other segments (i.e. the ones other than the j chosen ones), and so there’s nothing stopping some of them from having length \ge K.
Getting around this is exactly what inclusion-exclusion helps us do. Applying PIE to what we have, the number of valid configurations equals
This is the value of c_x, where M = N-x+1.
We only need to consider values of j such that N - j\cdot K \ge 0, which limits us to about j \le \frac{N}{K}.
However, applying this to each 0 \le x \le N still results in a complexity of \mathcal{O}(N^2 / K), which is still too slow for small K.
To further optimize the solution, we use the fact that we don’t actually care about the individual c_x values, but only the overall sum of c_x x! (N-x)!
Substituting in the expression derived for c_x, what we want to compute is
where M=N-x+1.
Swapping the order of summation, and expanding out M, we want
Looking at the expression inside the square brackets, observe that (by expanding the binomial coefficient):
Multiplying and dividing by (j\cdot K)! gets us
Thus, for a fixed j, we want to compute
The first two terms of this product are constants with respect to x, so we can take them out of the summation.
We thus want to find
This summation has a closed form, equal to exactly
(unless j = 0.)
Proof
While this can be shown algebraically, it’s always nice to have a combinatorial proof for the purposes of intuition - so that’s what we’ll do!
Let’s look at the summation we have.
It’s the sum of \binom{x}{jK}\binom{N-x+1}{j} over all values of x.
For a fixed x, this can be thought of as: picking jK items from x items, and then also picking j items from a separate collection of N-x+1 items.
There are thus N-x+1+x = N+1 items in total, across both collections.
We also add one additional item that is the separator between the collections, and enforce always picking this item - this is to allow us to easily sum over x.
Thus, we have N+2 items, out of which we’re picking jK+j+1 items. Here, the (jK+1)-th item picked will be our separator, corresponding to the value of x+1.
Of course, this means what we’re doing is equivalent to just picking j\cdot K + j + 1 items out of a collection of N+2, which can be done in \binom{N+2}{jK+j+1} ways - exactly what we claimed!
This works for all j \ge 1, and for j = 0 the summation comes out to be N+1, obviously.
With the above closed form for the inner summation, and treating j = 0 separately, the answer is hence
We only care about j \le \frac{N}{K} as noted earlier, and for a fixed j the expression can be evaluated in constant time so we’re done.
TIME COMPLEXITY:
\mathcal{O}(\frac{N}{K}) per testcase, after \mathcal{O}(\text{MAX\ \ N}) precomputation.
CODE:
Editorialist's code (C++)
// #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;
vector fac(200005, 1);
for (int i = 1; i < 200005; ++i) fac[i] = (fac[i-1] * 1ll * i) % mod;
auto inv = fac;
auto modpow = [&] (int a, int n) {
int r = 1;
while (n) {
if (n & 1) r = (1ll * r * a) % mod;
a = (1ll * a * a) % mod;
n /= 2;
}
return r;
};
for (auto &x : inv) x = modpow(x, mod-2);
auto C = [&] (int n, int r) {
if (n < r or r < 0) return 0;
int res = (1ll * fac[n] * inv[r]) % mod;
res = (1ll * res * inv[n-r]) % mod;
return res;
};
int t; cin >> t;
while (t--) {
int n, k; cin >> n >> k;
int ans = fac[n+1];
for (int j = 1; (k+1)*j <= n+1; ++j) {
int val = (1ll * C(n+2, (k+1)*j + 1) * fac[n - k*j]) % mod;
val = (1ll * val * fac[k*j]) % mod;
if (j & 1) ans += mod - val;
else ans += val;
ans %= mod;
}
cout << ans << '\n';
}
}