PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: q_ed
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Easy
PREREQUISITES:
Dynamic Programming
PROBLEM:
There are N cities. City i has color C_i and wealth W_i.
You are at city 1, which is guaranteed to have C_1 = W_1 = 0. Your initial color is 0.
You can do the following:
- Change your color to any value you like.
This has a cost of K coins. - Move from city i to city i+1.
You will gain W_{i+1} coins if your color equals C_{i+1}, and lose W_{i+1} coins otherwise. - The journey ends only when you reach city N.
Find the maximum possible number of coins you can have in the end.
EXPLANATION:
Observe that when we’re at city i, it’s only advantageous to change color to C_{i+1}.
This is because, if we change to some other color instead, we could just as well move to C_{i+1} first and then change to that color - in either case the cost is going to be the same.
This means that, if we’re currently at city i with color C_i (meaning we gained W_i coins from this city), the only thing that really matters is the previous city j \lt i where we gained coins (and so had a color of C_j there); because we’ll just lose coins for every city inbetween without changing color.
This dependence on only the previous “good” city brings up a somewhat natural dynamic programming solution.
Let’s define dp_i to be the maximum coins that can be obtained if we’re currently at city i with a color of C_i.
The transitions are simple in linear time: fix some j \lt i as the previous “good” city, and then the
coins obtained will equal dp_j + W_i - (W_{j+1} + W_{j+2} + \ldots + W_{i-1}), minus an additional cost of K if C_i \ne C_j.
This is because we’ll have dp_j coins at j, lose coins at every city from j+1 to i-1, (if needed) pay K coins to change color to C_i, and then gain W_i coins at city i.
This can be easily implemented in \mathcal{O}(N^2), since the quantity (W_{j+1} + W_{j+2} + \ldots + W_{i-1}) can be computed in constant time with the help of prefix sums.
The final answer will be the maximum of dp_i - (W_{i+1} + W_{i+2} + \ldots + W_N) across all i, since this corresponds to fixing the last “good” city and then we’ll lose coins at all future cities.
The constraints are too high for a quadratic solution to pass though, so we need to optimize this.
First, let’s rewrite the transition expression to make it a bit easier to deal with.
Define P_i = W_1 + W_2 + \ldots + W_i to be the i-th prefix sum of W.
Then, for a pair of indices j \lt i, we want to update dp_i with
- dp_j + W_i - (P_{i-1} - P_j) - K, if C_i \ne C_j
- dp_j + W_i - (P_{i-1} - P_j), if C_i = C_j
Let’s look at just the C_i = C_j case.
Regroup that expression to be (dp_j + P_j) + (W_i - P_{i-1}).
Observe that the first term depends purely on j, while the second depends purely on i.
For a fixed i, the second term is thus a constant; so we only need to care about maximizing the first term.
This is to be done across all j \lt i that share the same color, and since it depends only on j, we can just save the relevant value.
That is, define an array \text{best}, where \text{best}[x] stores the maximum value of (dp_j + P_j) across only indices j such that C_j = x.
Then, instead of having to iterate through all j, we can simply take \text{best}[C_i] + (W_i - P_{i-1}) for this case!
Optimizing the case of C_i \ne C_j is quite similar.
We can rewrite the expression to (dp_j + P_j) + (W_i - P_{i-1} - K).
Again, the second term is a constant, so we only need to maximize the first across all j \lt i; which can just be stored in a single variable.
We thus obtain the following algorithm:
- Build the prefix sum array P of W.
- Define an array \text{best}, initialized with all elements being some very small value (say, -10^{18}).
However, set \text{best}[0] = 0, since we start with C_1 = 0. - Also define a variable \text{globalBest}, initialized to 0.
- Then, for each i = 2, 3, \ldots, N in order:
- Set dp_i to \max(\text{globalBest} + (W_i - P_{i-1} - K), \text{best}[C_i] + (W_i - P_{i-1})).
This accounts for both cases of transitions. - Then, update \text{best}[C_i] and \text{globalBest} with the value of dp_i appropriately.
- Set dp_i to \max(\text{globalBest} + (W_i - P_{i-1} - K), \text{best}[C_i] + (W_i - P_{i-1})).
- Once all the values of dp are known, the final answer is the maximum of
dp_i - (P_N - P_i)
across all indices i.
TIME COMPLEXITY:
\mathcal{O}(N) per testcase.
CODE:
Editorialist's code (PyPy3)
for _ in range(int(input())):
n, k = map(int, input().split())
w = list(map(int, input().split()))
c = list(map(int, input().split()))
pref = w[:]
for i in range(1, n): pref[i] += pref[i-1]
dp = [0]*n
bestval = 0
bestcol = [-10**18]*n
bestcol[0] = 0
for i in range(1, n):
dp[i] = bestval - pref[i-1] - k + w[i]
dp[i] = max(dp[i], bestcol[c[i]] - pref[i-1] + w[i])
bestval = max(bestval, dp[i] + pref[i])
bestcol[c[i]] = max(bestcol[c[i]], dp[i] + pref[i])
ans = -sum(w)
for i in range(n):
ans = max(ans, dp[i] - (pref[-1] - pref[i]))
print(ans)