hydroshiba's blog

By hydroshiba, 7 weeks ago, 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! Comments (1)
 » nice solution