In this post we present a coding problem and our solution which is different from the author's. Our solution doesn't involve tree, and only a stack is used.
The original problem was phrased in psuedo code. In my wording it becomes:
A robot is to travel from station 0 to station 1, then to station 2, station 3, and so on, to station N - 1. It enters station 0 with energy 0 and no money. At each station i, the robot has two options: either it resets its energy to Pi ≥ 0, or it takes money Vi ≥ 0. As it enters next station, its energy deceases by 1. The only restriction is that its energy should never becomes negative. The goal is to maximize the total amount of money gathered. 1 ≤ N ≤ 500000, 0 ≤ Vi ≤ 109, 0 ≤ Pi ≤ 100000. It is guaranteed there's a way that energy is always non-negative.
A standard DP with state (station, energy) takes O[N × max_energy]. Let D(i, e) be the maximum amount of money to be gained when entering station i with energy e. Clearly
max money gained if reset energy at station i,
max money gained if take Vi at station i
Pi == 0 ? -1 : D(i + 1, Pi - 1),
e == 0 ? -1 : Vi + D(i + 1, e - 1)
As the complexity is too high, we have to try reducing it. Let's observe an example and see what we got.
N = 5
P: 3 0 1 3 1 V: 2 1 3 6 9
Our DP starts from i = 4. Because there's no concern of energy after entering the last station, D(4, e) are
9 9 9 9 9 9 9....
For i = 3, either we take energy 3 and enter i = 4 with energy 2 which will end up with 9 dollars, or we take 6 dollars if we enter with positive energy and end up with 6 + 9 = 15 dollars. So we're comparing these two rows:
9 9 9 9 9 9 9 ....15 15 15 15 15 15 ....
9 15 15 15 15 15 15 ...
For i = 2, either we take energy 1 and enter i = 3 with energy 0 which will end up with 9 dollars, or we take 3 dollars if we enter with positive energy and end up with 3 more dollars than i=3. So we're comparing these two rows:
9 9 9 9 9 9 9 ...12 18 18 18 18 18 ...
9 12 18 18 18 18 18 ...
As we see now, the pattern of update is to compare 2 rows. The first is obtained by expanding D(i + 1, Pi - 1) to a row. The second is by right shifting the row of D(i + 1, e) by 1 and add Vi. Adding Vi is expensive because the row can be long, so why not store Vi + Vi + 1 + ... + VN − 1 - earning in the table? This is exactly the amount of money we don't manage to get from station i to the end. This is the first observation.
The second observation is that the number of distinct values increases by at most 1, and is bounded by N − i. So what if we just store a vector or stack of segments with size no more than O(N)? One issue follows immediately: because of the right shift we have to update the end points of these segment in O(N) for each i. But since the right shift always occurs, what about storing end point - (N − i) as the coordinate of the segment? If a segment is not "interrupted", this value remains constant across stations and needs no update. In this way no O(N) update is necessary. In each iteration, the main task is to find the value of D(i + 1, Pi - 1) because it's not immediately available -- we have to perform O(log(N)) binary search of Pi − 1 on the end points of the segment. Once found, updating the stack takes amortized cost of O(1).
In brief, with these two coordinate transformations this problem is solved in O(N × log(N)) time and O(N) memory without tree structure. The code is available at https://github.com/paritystsai8/coding_problem/blob/master/robot.cpp