PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: iceknight1093
Tester: apoorv_me
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
Combinatorics
PROBLEM:
An array is called prefix-balanced if, for each of its prefixes, no two elements’ frequencies differ by more than 1.
You’re given N and M.
Find the number prefix-balanced arrays of length N, with elements from 1 to M.
EXPLANATION:
The very first thing we need to do is figure out how to characterize a prefix-balanced array.
Let S_A denote the set of elements present in A.
Then, observe that for any pair of elements x, y \in S_A:
- The second occurrence of x must be only after the first occurrence of y; otherwise at the point when x occurs twice the prefix will end up unbalanced.
- The third occurrence of x must be only after the second occurrence of y, for the same reason.
\vdots - The k-th occurrence of x must be only after the (k-1)-th occurrence of y.
In simpler terms, every element must occur once before any element occurs twice, every element must occur twice before any element occurs thrice, and so on.
Let K = |S_A|. Then, the above condition is only possible if:
- The first K elements are all distinct, i.e, a rearrangement of the elements of S_A.
- The next K elements are distinct, i.e, a rearrangement of S_A (although a different rearrangement can be chosen).
\vdots
Essentially, the array can be broken up into blocks of size K (except maybe the last block, which can have smaller size), and within each block the elements must be distinct.
We can use this characterization to count the number of prefix-balanced arrays.
- First, let’s fix K, the size of S_A. This is an integer between 1 and N.
- Next, let’s fix the K elements of S_A. For this, we can choose any K integers from \{1, 2, \ldots, M\}, giving us \binom{M}{K} choices in total.
- Recall that we broke up the array into several blocks of size K, and one with size (maybe) less than K.
Suppose there are B_K blocks of size K. Note that:- B_K will be equal to \left\lfloor \frac{N}{K} \right\rfloor.
- Within each of these K-sized blocks, any rearrangement of the elements of S_A can be chosen. There are K! rearrangements.
- With B_K blocks, the total number of possibilities is thus (K!)^{B_K}.
- That leaves us with the last block, of size \lt K (in particular, its size will equal (N\bmod K), the remainder when N is divided by K).
Let y_K = (N\bmod K). We can choose y elements out of the K we have, then arrange them in y! orders; for a total of \binom{K}{y}y! ways.
The final answer is obtained by just summing this up for all K \leq N, giving
where B_K = \left\lfloor \frac{N}{K} \right\rfloor and y_K = (N\bmod K).
Binomial coefficients can be computed in constant time if factorials and inverse factorials are precomputed, and (K!)^{B_K} can be computed in \mathcal{O}(\log B_K) (which is \mathcal{O}(\log N) using binary exponentiation.
This gives a solution in \mathcal{O}(N\log N) time.
TIME COMPLEXITY:
\mathcal{O}(N\log N) per testcase.
CODE:
Editorialist's code (Python)
mod = 998244353
N = 500005
fac = [1]*N
for i in range(1, N): fac[i] = fac[i-1] * i % mod
inv = fac[:]
inv[-1] = pow(fac[-1], mod-2, mod)
for i in reversed(range(N-1)): inv[i] = inv[i+1] * (i+1) % mod
def C(n, r):
if r < 0 or n < r: return 0
return fac[n] * inv[r] % mod * inv[n-r] % mod
for _ in range(int(input())):
n, m = map(int, input().split())
ans = 0
for i in range(1, n+1):
blocks = n//i
rem = n%i
choices = C(m, i) * pow(fac[i], blocks, mod) % mod
choices = choices * C(i, rem) % mod * fac[rem] % mod
ans += choices
print(ans % mod)
By the way can someone recommend me similar problems please?? where can I find them??