PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: mathmodel
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
Combinatorics
PROBLEM:
You’re given N and K.
Count the number of permutations P of [1, \ldots, N] such that P_i + P_{i+1} \not\equiv 0 \pmod K for each 1 \leq i \lt N.
In this version, K = 3.
EXPLANATION:
We only care about remainders modulo 3, so let’s reduce each integer in [1, N] to its remainder modulo 3 first.
Let c_0, c_1, c_2 be the number of values corresponding to remainders 0, 1, 2 respectively.
Let’s call an array A of length N valid if:
- A contains exactly c_0 copies of 0, c_1 copies of 1, c_2 copies of 2.
- A_i + A_{i+1} \not\equiv 0 \pmod 3 for each 1 \leq i \lt N.
There are then exactly c_0! \ \cdot\ c_1!\ \cdot\ c_2! ways to turn A into a permutation of [1, N].
So, we can just focus on computing the number of valid arrays A.
Given that we have 0 \leq A_i \leq 2, there are only really two ways to have A_i + A_{i+1} \equiv 0 \pmod 3.
These are to have two zeros adjacent to each other, or to have a 1 adjacent to a 2.
Let’s look at the zeros first.
We have c_0 of them, and they must be placed among N positions in such a way that no two are adjacent.
Suppose we’ve chosen a way to do this. How to place the ones and the twos?
Well, let’s look at two consecutive zeros, say at positions i and j.
Observe that in positions i+1, i+2, \ldots, j-1, we must place either only ones, or only twos - if we mix them, a 1 and a 2 will definitely be adjacent, which is not allowed.
So, we have several segments between the zeros, and what we want to do is decide for each segment whether to fill it with 1 or 2.
Further, the total length of segments filled with 1 must equal c_1 (note that this guarantees that the total length of segments filled with 2 will equal c_2, because c_0+c_1+c_2=N.)
Also, note that the number of available segments equals either c_0-1, c_0, or c_0+1, depending on whether a 0 is placed at the beginning and/or end of A or not (there are always c_0-1 segments between zeros, and maybe an additional one at the start/end).
However, it’s quite hard to actually count the number of ways to have such an assignment, because it depends greatly on the lengths of the segments themselves (and there are too many possible combinations of lengths to try and iterate through).
Instead, let’s approach the counting from the opposite direction - rather than fix the zeros and try to assign 1's and 2's to the resulting segments, we’ll instead fix the segments (in terms of their length and order), and then place 0's to separate them.
Now, if we fix just the number of segments that must contain 1's, it’s in fact quite easy to count the number of configurations.
After all, if we want x segments of ones, we’re just looking for the number of ways to write c_1 as the sum of x positive integers.
This is well-known to be \binom{c_1-1}{x-1}, from stars and bars.
Similarly, if we want y segments of 2's, there are \binom{c_2-1}{y-1} ways to obtain them.
If we have x segments of 1's, and y segments of 2's, there are \binom{x+y}{x} ways to arrange the segments in various orders.
Once we’ve fixed such an arrangement of segments, it in fact corresponds (almost) uniquely to a valid array A, by just placing a 0 between each adjacent pair of segments!
Or, more precisely,
- If there are c_0-1 segments, we obtain a valid array A by placing zeros between adjacent segments, as well as at the start and end of the array.
- If there are c_0+1 segments, we obtain a valid array A by placing zeros between adjacent segments only.
- If there are c_0 segments, we obtain exactly two different valid arrays: place zeros between adjacent segments, and then either a 0 at the beginning or at the end.
This gives us the following solution:
- Fix the number of segments, say to s.
- Fix the number of segments of 1's, say to x (0 \leq x \leq s).
- This then uniquely fixes y = s - x as the number of segments of 2's.
- There are then \binom{c_1-1}{x-1} \cdot \binom{c_2-1}{y-1} \cdot \binom{s}{x} ways to arrange x segments of 1's and y segments of 2's.
- If s = c_0, add twice this value to the answer; otherwise add it once (the reason is above).
We try a constant number of choices for s, and then each one is processed in \mathcal{O}(s) time.
Since s \leq N, this is linear time overall.
Don’t forget to multiply by c_0! \ \cdot\ c_1!\ \cdot\ c_2! in the end.
TIME COMPLEXITY:
\mathcal{O}(N) per testcase.
CODE:
Editorialist's code (PyPy3)
mod = 998244353
N = 500005
fac = [1]*N
for i in range(1, N): fac[i] = fac[i-1] * i % mod
inv = fac[:]
for i in range(N): inv[i] = pow(fac[i], mod-2, mod)
def C(n, r):
if n < r or r < 0: return 0
return fac[n] * inv[r] * inv[n-r] % mod
def sums(n, k):
# n positive integers summing to k
if k < n: return 0
if n == k: return 1
return C(k-1, n-1)
for _ in range(int(input())):
n, k = map(int, input().split())
x, y, z = n//3, n//3, n//3
if n%3 > 0: y += 1
if n%3 == 2: z += 1
ans = 0
for m in [x-1, x, x, x+1]:
for c in range(m+1):
ans += C(m, c) * sums(c, y) * sums(m-c, z) % mod
print(ans * fac[x] * fac[y] * fac[z] % mod)