Alternate solution to AtCoder Educational DP Problem M

#	User	Rating
1	tourist	4009
2	jiangly	3823
3	Benq	3738
4	Radewoosh	3633
5	jqdai0815	3620
6	orzdevinwang	3529
7	ecnerwala	3446
8	Um_nik	3396
9	ksun48	3390
10	gamegame	3386

#	User	Contrib.
1	cry	167
2	Um_nik	163
3	maomao90	162
3	atcoder_official	162
5	adamant	159
6	-is-this-fft-	158
7	awoo	157
8	TheScrasse	154
9	Dominater069	153
9	nor	153

I've just came up with an optimization for problem M of AtCoder Educational DP contest (link) which only uses mathematical deductions and needs no use of data structures. If you want a brief solution that uses prefix sum, there's already a pretty great blog here.

First, let $$$dp[i][k]$$$ be the number of ways of distributing candies to the first $$$i$$$ kids and have $$$k$$$ candies left. We can try giving the $$$i_{th}$$$ kid $$$0$$$, $$$1$$$, $$$2$$$, ... candies and the number of candies left should be $$$k$$$, $$$k - 1$$$, $$$k - 2$$$, ...

This leads to the recurring formula:

$$$ \large dp[i][k] = \sum\limits_{j = 0}^{a_i} dp[i - 1][k - j] $$$

This formula has $$$N * K$$$ states, and since $$$a_i$$$ can be as large as $$$K$$$, each transition costs $$$O(K)$$$ which makes the total complexity $$$O(N * K^2)$$$. Clearly with $$$K \leq 10^5$$$, this won't get us anywhere, so let's optimize it.

For a moment let's forget $$$dp[i][k]$$$ and examine $$$dp[i][k - 1]$$$ instead. Based on the recurring formula we have:
ㅤ

$$$ \large dp[i][k - 1] = \sum\limits_{j = 0}^{a_i} dp[i - 1][k - 1 - j] = \sum\limits_{j = 1}^{a_i + 1} dp[i - 1][k - j] $$$

Let's do some magic to this formula. By expanding this a bit, we have:
ㅤ

$$$ \begin{align*} &\sum\limits_{j = 1}^{a_i + 1} dp[i - 1][k - j]\\ &=\sum\limits_{j = 1}^{a_i} dp[i - 1][k - j] + \color{red}{dp[i - 1][k - (a_i + 1)]}\\ &=\sum\limits_{j = 1}^{a_i} dp[i - 1][k - j] + \color{cornflowerblue}{dp[i - 1][k - 0]} - \color{cornflowerblue}{dp[i - 1][k - 0]} + \color{red}{dp[i - 1][k - (a_i + 1)]}\\ &=\sum\limits_{j = \color{cornflowerblue}{0}}^{a_i} dp[i - 1][k - j] - \color{cornflowerblue}{dp[i - 1][k]} + \color{red}{dp[i - 1][k - (a_i + 1)]}\\ \end{align*} $$$

Notice how I merged the term $$$\color{cornflowerblue}{dp[i - 1][k - 0]}$$$ into the summation. Now, some of you will realize that $$$\sum_{j = 0}^{a_i} dp[i - 1][k - j]$$$ is also conveniently $$$dp[i][k]$$$, so we can rewrite the formula as:
ㅤ

$$$ \begin{align*} &dp[i][k - 1] = \color{purple}{dp[i][k]} - \color{cornflowerblue}{dp[i - 1][k]} + \color{red}{dp[i - 1][k - (a_i + 1)]}\\ &\Leftrightarrow - \color{purple}{dp[i][k]} = - dp[i][k - 1] - \color{cornflowerblue}{dp[i - 1][k]} + \color{red}{dp[i - 1][k - (a_i + 1)]}\\ &\Leftrightarrow \color{purple}{dp[i][k]} = dp[i][k - 1] + \color{cornflowerblue}{dp[i - 1][k]} - \color{red}{dp[i - 1][k - (a_i + 1)]}\\ \end{align*} $$$

And behold, the formula which only costs $$$O(1)$$$ for transition! Just by using some rearrangements, we have a totally new recurring formula that allows us to solve the problem in $$$O(N * K)$$$, and this is just beautiful.

Thanks for reading my blog. Upvote if you like it, downvote if you don't and have a good day!

hydroshiba's blog