This blog discusses a problem from the I2024 ICPC Chennai Regionals. My solution differs from the author's, so I decided to write a blog (and maybe gain some contribution). The problem I shall discuss is Problem F (Easy Counting Problem) from the 2024 ICPC Chennai Regional Problem Set.
Special thanks to satyam343 for motivating this blog and UnexpectedValue for helping me write it. Also Thanks to temporary1, ALnQ417 and satyam343 for proof reading the blog.
Problem Statement
We delete an element $$$A[i]$$$ from array $$$A$$$ if $$$A[i] = i$$$, and $$$i$$$ is the rightmost such index if there are multiple such indices. $$$\newline$$$
Constraints $$$ \newline $$$ $$$ N, M \leq 2 \cdot 10^5 , \space |N - M| \leq 20 \space and \space 1 \leq A[i] \leq M \newline $$$ Note : Array is one-indexed
We need to compute the sum of $$$\text{cost}(A)$$$ over all arrays $$$A$$$ such that $$$A[i] \leq M$$$.
Dynamic Programming Approach
Let's define: $$$ \text{DP}[i][j] = \text{In the prefix of } i \text{ elements, we can only delete } j \text{ elements.} $$$
$$$\color{red}{\text{For}}$$$ $$$\color{red}{i \leq M}$$$: $$$DP[i][j] = DP[i - 1][j] \cdot (M - 1 - j) + DP[i - 1][j - 1] \cdot j$$$
$$$\color{red}{\text{For}}$$$ $$$\color{red}{i > M}$$$, there is a similar DP but with more restricted states. Since $$$|N - M| \leq 20$$$, if we find $$$DP[M][j]$$$, we can solve the remaining part in $$$|N - M| \times N$$$ complexity.
Improved Polynomial Approach
If we approach the question using the naive dynamic programming approach as above, the complexity will be $$$\mathcal{O}(N^{2})$$$. We will now try to improve our complexity by trying to find a polynomial for $$$F_j$$$ where coefficient of $$$x^i$$$ in $$$F_j$$$ is $$$\text{DP[i][j]}$$$.
Define $$$m = M - 1$$$ for simplicity.
$$$F_j = x \cdot F_j \cdot (m - j) + x \cdot F_{j - 1} \cdot j$$$
Rewriting: $$$F_j = \frac{jx \cdot F_{j - 1}}{1 - (m - j)x}$$$, with $$$F_0$$$ as shown below.
This follows from: $$$DP[i][0] = DP[i - 1][0] * m$$$ which gives coefficients $$${1, m, m^2, m^3, \dots}$$$, forming a geometric progression with summation $$$F_0 = \frac{1}{1 - mx}$$$.
Now, $$$F_j = \frac{j! \cdot x^j}{(1 - mx)(1 - (m-1)x) \dots (1 - (m-j)x)}$$$
This is inefficient, so we use:
This follows the general formula: $$$F_j = \sum_{k=0}^{j} (-1)^k \cdot \binom{j}{k} \cdot \frac{1}{1 - (m - k)x}$$$, which can be proved using the Principle of Mathematical Induction. $$$\newline$$$
So we can have $$$\text{dp[n][j]} = \sum \left((-1)^k \cdot \binom{j}{k} \cdot (m - k)^n\right)$$$. Now let's try to find values of $$$\text{dp[n][j]}$$$ for all $$$j$$$, which means now $$$n$$$ is constant. Let's try to break it down as: $$$\newline$$$
We got lucky in the fact that $$$\text{mod}$$$ is $$$998244353$$$, and I think many might have seen it as the multiplication of 2 polynomials already.
Let's say the polynomials
are enough. Now after multiplying P and Q we get polynomial $$$F$$$ whose coefficients are $$$\text{DP[N][j]}$$$.
Let's transition from $$$\text{DP[i]}$$$ to $$$\text{DP[i + 1]}$$$ if $$$i \geq M$$$. Let $$$A[j] =$$$ number of ways we can select a value here such that the value at the higher level changes and $$$B[j]$$$ such that it does not change for $$$\text{DP[i][j]}$$$. $$$\newline$$$
Let's get the value for $$$A[i]$$$ and $$$B[i]:$$$
- We have $$$A[j] = (M - (i - j))$$$ as we can remove $$$j$$$ elements up until now, so $$$\text{min value of } (i - j) + 1$$$ will change the array, and all values until $$$i + 1$$$ will change. But as $$$M < i + 1$$$, $$$(i + 1)$$$ cannot be the largest element, so $$$(M - (i - j))$$$ is the amount all remaining values will not change the number of deleted elements.
- and $$$B[j] = M - A[j + 1]$$$.
Using this transition to go from $$$\text{DP[i]}$$$ to $$$\text{DP[i + 1]}$$$, we can solve the problem in $$$\mathcal{O}((N - M) \cdot N + Nlog(N))$$$ Solution.
