Ah ! It has been a DP Sunday. I have done a good amount of DP earlier but as far as I recall, had never gone above 2D DP problems. So, while moving ahead with my A2OJ streak, I came across D. Caesar's Legions .
Before discussing the problem, I would like to mention about an awesome tutorial I found while figuring out the solution: https://noi.ph/training/weekly/week3.pdf . It goes through this problem in a beautiful fashion.
In the post, I shall be going through 2 different solutions with different running complexity. So, let's begin.
The problem asks you to find the number of sequences made up of n1 footmen and n2 horsemen such that no more than k1 footman are together and not more than k2 horsemen are together.
Let us construct a 4D DP array where DP[i][j][k][l] refer to number of ways given i footmen, j horsemen, such that we can still add k footmen or l horsemen to the back of the sequence.
Now the transitions are as follows:
Base Case: When i = 0 and j = 0, clearly we can have only 1 sequence (an empty sequence !). So, DP[k][l] = 1.
Now there are 2 possible cases:
- We add a footmen to the back. Clearly, this is possible only if i > 0 (we have atleast one footmen) and k > 0 (we can atleast add a footmen to the back (there is no suffix of present sequence with k1 footmen)). Thus we add DP[i - 1][j][k - 1][l] .
- Similarly, for horsemen, we add DP[i][j - 1][k][l - 1] to the answer if j > 0 and l >0.
There are n1 * n2 * k1 * k2 possible states and calculating the value for each takes O( 1 ) time. Thus, the complexity is O( n1 * n2 * k1 * k2 ).
Thanks to guys on CodeForces Discord, after struggling in understanding the official editorial, I was able to get it and finally was able to reach an O( n1 * n2 * max(k1, k2) ) solution.
Let us create a 3D DP where DP[i][j][k] represents number of ways to construct the sequence with i of footmen, j of horsemen, and k = 0 if we add a footmen at this step and k = 1 if we add a horsemen at this step.
Base case is similar to Solution 1. If i == 0 and j == 0, DP[i][j][k] = 1.
Now let us see the possible transition states.
- DP[i][j] : This means that the last element we are going to add is going to be a footmen. We know that we can add a maximum of k1 footmen together. Thus before the suffix of footmen, there must be a horsemen. As we can add upto k1 horsemen, we do dp[i][j] += dp[i - together][j], where together is the number of footmen that we add at the end, which is [1, min(i, k1)]. And why  ? This is because just before this suffix of footmen, there must be a horsemen !.
- In a similar fashion, we can get the recursion relation for DP[i][j].
Total number of states is clearly n1 * n2. For each state, we take min(i, k1) or min(i, k2) time. Thus for each state, at max we take time equivalent to max(k1, k2). This implies that overall complexity is O(n1 * n2 * max(k1, k2) ).
PS: I think I should start writing a daily unofficial editorial of problems I solved that day. Will help me recall later and also improve my writing skills :P .
This article was updated on April 14, 2019