Contest Division 1
Contest Division 2
Contest Division 3
Contest Division 4

Setter: Satyam Kumar Roy
Testers: Nishank Suresh and Abhinav sharma
Editorialist: Satyam Kumar Roy






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.


  • 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


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.


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.


The time complexity is O( \frac{MAX \cdot MAX}{b} + q \cdot b), where MAX=10^5 and b=500.


Setter and Editorialist’s Solution

Tester 1’s Solution

Tester 2’s Solution

Feel free to share your approach. Suggestions are welcomed as always. :slight_smile:


Now on multiplying by p^r on both sides, we get

some r in the blow formulas should be p.

\binom {n-1} {r-1} \cdot p^{r-1} \cdot p

p \cdot c(n - 1, r - 1)

I think there is a typo in the equation…
I think it should be p*c(n-1,r-1)+c(n-1,r) and any subsequent derived equations should also be changed.

Thank You!
I have edited the editorial.