Invite the Heroes

Here, we have to do three things optimally.

1. Choose the targeted city, where we want to invite the heroes.

2. Choose exactly one hero from each type.

3. Distribute the energy drinks among the invited heroes.

The number of cities can be at most $5000$. Thus, it is obvious to invite the heroes to each city and find the optimal result in $O(n)$ for that city. The final complexity will be $O(n^2)$.

Let’s group the heroes by their types. So, there are $k$ groups and from each group, we have to invite exactly one hero. It can be done using dynamic programming (0-1 knapsack).

Suppose, we have invited the heroes of index $x_1,~x_2,~x_3,~…,~x_k$ and the targeted city is $z$ where $x_1 < x_2 < x_3 < … < x_m < z \leq x_{m+1} < … < x_k$. The power of invited heroes after gathering without drinking energy drinks are $y_1,~y_2,~y_3,~…,~y_k$. $(y_i = max(0, ~p_{x_i}-|x_i - z|)$.

Now, the cities of the first $m$ invited heroes are on the left side of the targeted city. We want to distribute the energy drinks optimally among these $m$ heroes. The leftmost hero of index $x_1$ can reach index $x_2$ without losing any power by drinking energy drinks. So there are $2$ heroes at index $x_2$ with power $y_1^\prime,$ and $y_2$ respectively where $y_1^\prime = max(0, ~p_{x_1} - |x_2 - z|)~$($y_1^\prime$ is the power of the hero of city $x_1$ after reaching at city $z$ without drinking energy drink and starting from city $x_2$).

If both $y_1^\prime$ and $y_2$ are greater than $0$, then we can take anyone between the heroes of index $x_1$ and $x_2$ at city $x_3$ giving all the energy drinks between city $x_2$ and $x_3$. Both of the results are optimal.

If $y_2$ is $0$, then we have to take the hero of city $x_1$ at city $x_3$ giving the energy drinks between city $x_2$ and $x_3$.

If $y_1^\prime$ is $0$, then we must take the hero of index $x_2$ at $x_3$ giving the energy drinks between city $x_2$ and $x_3$. This is because taking $x_2$ will not contribute anything and the energy drinks between city $x_1$ and $x_2$ have no effects on the optimal result in this case.

So we get the optimal result by moving the hero of city $x_1$ or $x_2$ with initial power (drinking energy drinks while moving will not decrease his power and preserve the initial power). It is always better to assign energy drinks to exactly one hero to the left of $z$ and to exactly one hero to the right of $z$.

Then our problem comes down to a modified 0-1 knapsack, where with taking the heroes, we also need to track if we have already assigned the drinks to a left hero or a right hero. Complexity $O(n^2)$.