Author: Stepan Filippov
Tester: Mark Mikhno
Editorialist: Bhuvnesh Jain
You are given a list of N dragons each having some power A[i]. There is a prince with power W. The princess decides to sit at position x, 1 <= x <= N. The prince starts from any position towards the princess and loses power equal to A[i] when he is at index i. If dies at any moment are power becomes non-positive. She wants to arrange the dragons in such a way such that the number of positions from which prince dies before reaching her is maximised. She want to calculate this value for all x, 1 <= x <= N.
The first to note is that we would always like to place stronger dragons close to the princess as it will increase the number of positions from which prince would die. This equivalently means, we would like to have the minimum number of dragons such that their sum of powers exceeds W. For this, we have two conditions, either we place dragons on both sides such that their sum of powers exceeds W, meaning there are places on both sides such that the princess will die. In another case, we place dragons on one side only such that their sum of powers exceeds W so that princess doesn’t die from one side but dies from some places on another side.
The other thing to note is that the largest dragon will always stay at the place the princess is staying. This is consistent with the fact that the princess wants stronger dragons to stay close to her so that prince dies from the maximum number of starting positions. So, when considering dragons to be placed on both sides to defeat the prince, this dragon will be counted twice and this needs to be taken care in implementation.
Now, in both cases we have the following reduced problems:
For one side case, we need to find the minimum number of dragons such that the sum of their powers is greater than or equal to W.
For both side case, we need to find minimum number of dragons such that for index i, there is partition of size less than or equal to (i - 1) with sum of powers greater than or equal to W and other disjoint partition of size less than or equal to (N - i) such that sum of powers is greater than or equal to W.
The answer for each index is basically the maximum answer we can obtain from both cases.
For the first case, the answer can be found greedily by simply placing the dragons in decreasing order of power and finding the first index such that sum of powers till now is greater than or equal to W. Let us denote this value by z.
For the second case, let us simply the problem further. For this, let us think if the size of partitions really matter, (i.e. do we care about the index i for which the calculation is being done). So, let us assume the optimal partitions of dragons into 2 disjoint sets such that the sum of both is greater than or equal to W have sizes x and y. It is evident that z <= x and z <= y. For now, assume left size size is less than right side size i.e. (i - 1) < (N - i). Also, x <= y.
If x <= (i - 1) and (y <= (N - i)), then everything is good. Assume x > (i - 1), o dragons can’t protect the princess from the left side and y dragons protect the princess from the right side. In this case, it is better to choose z dragons only to protect the princess from one side as z <= y and it makes our answer better. So, the size of x and y don’t matter in this case. Again if x > (i - 1) and y > (N - i), it means princess can never be protected from both sides and we need to consider the answer for one side case only, meaning sizes of x and y don’t matter again.
We can prove it for other cases like (i - 1) > (N - i) similarly. Thus, we conclude size of partition don’t matter and the reduced problem is therefore
Find the minimum sum of x and y such that there is a disjoint partition of size x and y such that sum of powers of the dragon in both partitions is greater than W
Let us first calculate the minimum number of dragons required to obtain the sum of powers as greater than or equal to U. This is the same as the reduced problem 1. We will calculate this value for every U, 1 <= U <= 3 * W, as till will be required later.
For this, we will use dynmaic programming solution. Let dp[i] = 1 if there exists a partiton of dragons such that their sum of powers is exactly i. Otherwise, dp[i] = 0.
We will process all equal powers simultaneously. Let the current power be u. We consider the sum of partitions A and B. Let A = j + k * u, for valid j and k. We need A >= W, meaning for given j, we can find the minimum number of dragons required for A >= W. The only other condition is B >= W. So, we need to find the minimum number of dragons with the sum of powers as (A + W). This can be done using the previous pre-computation we did. We can thus update our answer, only if dp[j] is set. This is because, if dp[j] is not set, we can’t obtain partition for A.
Now, we need to update the possible dp[j] values as we can use some number u. For this, we consider another dynamic programming, gp[v] denoting minimum number of u required to set dp[v]. The trivial case is when dp[v] is already set, meaning gp[v] = 0. For other cases, we try to extend the solution for (v - u) by 1, meaning extend the sequence by u.
Once, gp[v] is calcualted, we update the dp[v] as follows: dp[v] = dp[v] or (gp[v] <= m) where m is the number of dragons with power u. This works because:
- If dp[v] was already set, we ignore it.
- If gp[v] > m, we don’t have required number of dragons with power u.
- If v < u, then we should have gp[v] = inf, i.e. a large value.
- If v == u, then we have gp[v] = 1 and we can obtain v using a dragon of power u.
- If v > u, we can obtain v only if we have suitable number of dragons with power u i.e. gp[v] <= m. The proof lies in the way gp[v] is constructed and extended using dragons of power u.
For details, I request you to go through the commented solution of the author below. The case analysis is explained in line with the editorial.
For time complexity analysis:
- O(W) for calculating min_pref array in solution.
- O(min(N, sqrt(W)) * W) for doing the dynamic programming solution for both sides.
- O(N) for calculating the final answer.
For case 2, the proof (from author) is as follows:
We process all equal values u in O(W). Note that any 1.5 * sqrt(W) different positive numbers have sum >= W, so if we take 3 * sqrt(W) different positive numbers we can split the into two sets with sums >= W. My solution breaks the loop when p >= ans. If we’ve processed 3 * sqrt(W) different numbers then we’ve found the partition into two sets containing at most 3 * sqrt(W) different values. Let q be the index of (3 * sqrt(W)+1)th value. When (p = q) variable (ans < q), but all next possible partitions will contain q, so they won’t be optimal. So we can safely break our loop.
So the time complexity is O(min(N, sqrt(W)) * W).
If your solution was somewhat different, feel free to share in the comments section below.
O(N + W + min(N, sqrt(W)) * W)
O(N + W)
Author’s solution can be found here.
Tester’s solution can be found here.