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_{i1}. 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_{i1} (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_{i1} (2 \leq i \leq n)
Now g(n,m,k)= {{n1} \choose k} \cdot m \cdot (m1)^{n1k}
How to get this expression?
First select k indices i (2 \leq i \leq n) such that a_i=a_{i1}. We have n1 options(from 2 to n) for that, and we need to select k indices from them. So we get {{n1} \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_{i1} and we have already fixed the value of a_{i1}.
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_{i1}. So, we have m1 options for a_i.
Why?
We can have any integer from 1 to m, except a_{i1}, thus m1 options.
We have n1k such indices i. Hence we have m1 options for n1k indices. We get (m1)^{n1k} for this part.
To sum up, we have

m options for 1 index

1 option for k indices

m1 options for n1k indices
Thus, we get g(n,m,k)= {{n1} \choose k} \cdot m \cdot (m1)^{n1k}
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} {{n1} \choose i} \cdot m \cdot (m1)^{n1i}
f(n,m,k)= m \cdot (m1)^{n1} \sum_{i=0}^{k} {{n1} \choose i} \frac{1}{(m1)^i}
Let p=\frac{1}{(m1)}.
So f(n,m,k)= m \cdot (m1)^{n1} \sum_{i=0}^{k} {{n1} \choose i} p^i
So if we can calculate \sum_{i=0}^{k} {{n1} \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} = {{n1} \choose {r1}}+{{n1} \choose r} (n, r \geq 1)
Now on multiplying by p^r on both sides, we get
{n \choose r} \cdot p^r = {{n1} \choose {r1}} \cdot p^{r1} \cdot r+{{n1} \choose r} \cdot p^r
It is equivalent to
c(n,r)= p \cdot c(n1,r1) + c(n1,r)
Thus s(n,r)= p \cdot s(n1,r1) + s(n1,r)
We know that s(n1,r)=s(n1,r1)+c(n1,r)
Thus on substituting s(n1,r), we get
s(n,r)= p \cdot s(n1,r1) + s(n1,r1)+c(n1,r)
Finally s(n,r)= (p+1) \cdot s(n1,r1) + c(n1,r)
After precomputing factorials and their inverses, we can caculate c(n1,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(n1,r1) + c(n1,r).
So, to know s(n,r), we need s(n1,r1). Similarly to know s(n1,r1), we need s(n2,r2), 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(nk,rk).
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.