PROBLEM LINK:
Contest Division 1
Contest Division 2
Contest Division 3
Contest Division 4
Practice
Setter: Satyam Kumar Roy
Testers: Nishank Suresh and Abhinav sharma
Editorialist: Satyam Kumar Roy
DIFFICULTY
2787
PREREQUISITES
Combinatorics
PROBLEM
Suppose f(X, Y, Z) denotes the number of arrays A of length X such that:
- 1 \leq A_i \leq Y, for all (1 \leq i \leq X);
- There are at most Z values of i (2 \leq i \leq X) such that A_i=A_{i-1}. In other words, for at most Z indices, the element at a given index is equal to the previous element.
You are given an integer M and Q queries. The i^{th} query is of the form:
- N_i K_i: Given integers N_i and K_i, find the value f(N_i,M,K_i).
Since output might be too large, print it modulo 998244353.
QUICK EXPLANATION
-
Suppose g(n,m,k) gives number of arrays a of length n such that
-
1 \leq a_i \leq m
-
there are exactly k indices i such that a_i=a_{i-1} (2 \leq i \leq n)
-
-
We can find expression for g(n,m,k)
-
We can express f(n,m,k) in terms of g(n,m,i) (0 \leq i \leq k)
-
Now on smartly precomputing some values, we can find f(n,m,k) quickly
EXPLANATION
Suppose g(n,m,k) gives number of arrays a of length n such that
-
1 \leq a_i \leq m
-
there are exactly k indices i such that a_i=a_{i-1} (2 \leq i \leq n)
Now g(n,m,k)= {{n-1} \choose k} \cdot m \cdot (m-1)^{n-1-k}
How to get this expression?
First select k indices i (2 \leq i \leq n) such that a_i=a_{i-1}. We have n-1 options(from 2 to n) for that, and we need to select k indices from them. So we get {{n-1} \choose k} for this part.
Say selected indices lies in set S.
Now that we have selected k indices, let us start putting values.
We can have any value for a_1. So we have m options for a_1.
Now let us fix values of indices from 2 to n in this order. We will first fix a_2, then a_3 and so on.
Start from i=2, and move towards right.
If i lies in set S, a_i is already fixed, as a_i=a_{i-1} and we have already fixed the value of a_{i-1}.
So we have 1 option for all such i. We have 1 option for k indices.
If i does not lie in set S, this means a_i \neq a_{i-1}. So, we have m-1 options for a_i.
Why?
We can have any integer from 1 to m, except a_{i-1}, thus m-1 options.
We have n-1-k such indices i. Hence we have m-1 options for n-1-k indices. We get (m-1)^{n-1-k} for this part.
To sum up, we have
-
m options for 1 index
-
1 option for k indices
-
m-1 options for n-1-k indices
Thus, we get g(n,m,k)= {{n-1} \choose k} \cdot m \cdot (m-1)^{n-1-k}
In the problem, we have to find f(n,m,k) for many queries.
Note that f(n,m,k)=\sum_{i=0}^{k} g(n,m,i).
f(n,m,k)=\sum_{i=0}^{k} {{n-1} \choose i} \cdot m \cdot (m-1)^{n-1-i}
f(n,m,k)= m \cdot (m-1)^{n-1} \sum_{i=0}^{k} {{n-1} \choose i} \frac{1}{(m-1)^i}
Let p=\frac{1}{(m-1)}.
So f(n,m,k)= m \cdot (m-1)^{n-1} \sum_{i=0}^{k} {{n-1} \choose i} p^i
So if we can calculate \sum_{i=0}^{k} {{n-1} \choose i} p^i for each query, we are done.
Now let us see how to calculate this quickly.
Suppose c(i,j)={i \choose j} p^j (p is same as the one used above)
Suppose s(i,j)=\sum_{k=0}^{j} c(i,k)
It is well known that {n \choose r} = {{n-1} \choose {r-1}}+{{n-1} \choose r} (n, r \geq 1)
Now on multiplying by p^r on both sides, we get
{n \choose r} \cdot p^r = {{n-1} \choose {r-1}} \cdot p^{r-1} \cdot r+{{n-1} \choose r} \cdot p^r
It is equivalent to
c(n,r)= p \cdot c(n-1,r-1) + c(n-1,r)
Thus s(n,r)= p \cdot s(n-1,r-1) + s(n-1,r)
We know that s(n-1,r)=s(n-1,r-1)+c(n-1,r)
Thus on substituting s(n-1,r), we get
s(n,r)= p \cdot s(n-1,r-1) + s(n-1,r-1)+c(n-1,r)
Finally s(n,r)= (p+1) \cdot s(n-1,r-1) + c(n-1,r)
After precomputing factorials and their inverses, we can caculate c(n-1,r) in O(1).
Here comes the interesting part!
If we precompute s(n,i) (0 \leq i \leq n), if and only if n is divisible by 500 (which can be done in O(\frac{MAX}{500} \cdot MAX)), we can find s(n,r) in atmost 500 iterations. Here MAX=10^5.
We can find s(n,r) this way.
We have s(n,r)= (p+1) \cdot s(n-1,r-1) + c(n-1,r).
So, to know s(n,r), we need s(n-1,r-1). Similarly to know s(n-1,r-1), we need s(n-2,r-2), and so on.
Now we know that s(n,0)=1, and we also know s(x,y) if x is divisible by 500.
For pair (n,r), assume k is smallest non negative integer such that we already know s(n-k,r-k).
It is easy to see that 0 \leq k < 500.
Hence, we can can calculate answer for each query in atmost 500 iterations.
TIME COMPLEXITY
The time complexity is O( \frac{MAX \cdot MAX}{b} + q \cdot b), where MAX=10^5 and b=500.
SOLUTIONS
Setter and Editorialist’s Solution
Feel free to share your approach. Suggestions are welcomed as always.