BURGMOV - Editorial

editorial
, ,
1

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: raysh07
Tester: sushil2006
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Dynamic Programming

PROBLEM:

There are N houses in a row. The i-th house has A_i coins in it.
There is a burglar initially at house S.
In one day, the burglar can either move one step left/right, do nothing, or rob the current house and take all its coins (if it hasn’t been robbed from previously.)

For each S = 1, 2, 3, \ldots, N, find the maximum number of coins the burglar can obtain in K days.

EXPLANATION:

Let’s solve for a fixed start S.

The burglar must start at S, and can only move one step at a time.
So, if L denotes the leftmost house ever visited by the burglar and R denotes the rightmost house, all the houses in the range [L, R] will be visited by the burglar at least once.
For each house, we can thus decide whether we want to rob it or not.

Further, if L and R are fixed, it’s easy to see that the robber’s movement will be either:

  • Start at S, repeatedly move left till L is reached, then repeatedly move right till R; or
  • Start at S, repeatedly move right till R is reached, then repeatedly move left till L.

Clearly, it’s best to choose whichever one requires fewer movements (which will be R-L + \min(S-L, R-S)), so that the number of houses that can be robbed is maximized.
Since the number of houses that can be robbed is a constant (subtract the movement cost from K), say C, it’s clearly to optimal to only rob the C houses with the largest number of coins in the range.

This gives a solution in approximately quadratic time (maybe with a log factor) where we try all possibilities of (L, R).

However, this was only for a single starting location S, and running it for each one is too expensive for the given constraints.

To extend the solution, observe that we can actually split the movements in [L, R] into two independent parts, each with S as one endpoint.
Specifically, if we move left and then right, we can think of this as:

  • Start at S, then move left several times and then return to S.
  • From S, then move right several times, but without returning to S.

Further, rather than think about the exact borders (L, R), we only really need to ensure that the total cost on both sides (of both movement and robbing) doesn’t exceed K.

With this in mind, let’s define a function \text{return}_L(S, x) to be the maximum possible number of coins the robber can each such that:

  • The starting location is S.
  • The robber moves left several times, robs a few houses along the way, and then returns to S.
  • The total number of moves made is exactly x.

With this definition, transitions are quite easy.
There are only two options:

  1. Don’t rob house S.
    Here, we need to spend one move moving left to S-1, then we’ll move left from S-1 and rob some houses before returning to S-1, and then spend one more move to return to S.
    This means there are x-2 moves made from S-1, so by definition we obtain \text{return}_L(S-1, x-2) coins as the maximum possible.
  2. Alternately, we could rob from house S.
    The analysis is the same as above, just that we spend one move to obtain A_S coins instead.
    So, the best we can do here is \text{return}_L(S-1, x-3) + A_S.

\text{return}_L(S, x) equals the maximum of these two values.
There are \mathcal{O}(NK) states and constant-time transitions from each, so these can all be computed quickly.

Similarly, we can compute the following values in \mathcal{O}(NK) time:

  • \text{return}_R(S, x), the maximum number of coins that can be obtained by starting at S, moving right and returning to S, with a total cost of x.
    • The transitions are basically the same as \text{return}_L, just using S+1 instead of S-1.
  • \text{straight}_L(S, x), the maximum number of coins that can be obtained by starting at S, moving left and not turning back, with a total cost of x.
    • The transitions here are either \text{straight}_L(S-1, x-1) or \text{straight}_L(S-1, x-2) + A_S, depending on whether we steal from S or not.
  • \text{straight}_R(S, x), the same thing as \text{straight}_L but for moving right instead.

With all these known, we can now solve for a fixed starting point S in \mathcal{O}(K) time.

There are two possibilities: either we move left, return to S, and move right; or do the opposite.
Suppose we move left first.
If we spend a cost of x on the left, then,

  • The best we can do is \text{return}_L(S, x) on the left side.
  • That leaves a cost of K-x on the right.
  • One of them is needed to move from S to S+1, so that leaves K-x-1.
  • Once at S+1, we don’t need to return.
    So, the best we can do on the right is \text{straight}_R(S+1, K-x-1).
    (More precisely, it is the maximum of \text{straight}_R(S+1, y) across all y \le K-x-1; so just take the prefix maxima first.)

Note that after fixing x the overall profit can be computed in \mathcal{O}(1), so we can simply try all 0 \le x \le K and take the best one.

Similarly, in \mathcal{O}(K) time we can compute the best possible cost of starting at S, moving right and returning to S, and then moving left.
Report the best among all these answers.

We used \mathcal{O}(NK) precomputation and afterwards spend \mathcal{O}(K) time per index, so this is easily fast enough.

TIME COMPLEXITY:

\mathcal{O}(NK) per testcase.

CODE:

Editorialist's code (C++) 
// #include <bits/allocator.h>
// #pragma GCC optimize("O3,unroll-loops")
// #pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")
#include "bits/stdc++.h"
using namespace std;
using ll = long long int;
mt19937_64 RNG(chrono::high_resolution_clock::now().time_since_epoch().count());

int main()
{
    ios::sync_with_stdio(false); cin.tie(0);

    int t; cin >> t;
    while (t--) {
        int n, k; cin >> n >> k;
        vector a(n+2, 0);
        for (int i = 1; i <= n; ++i) cin >> a[i];

        vector dpL1(n+2, vector(k+1, 0));
        auto dpL2 = dpL1;
        auto dpR1 = dpL1;
        auto dpR2 = dpL1;

        for (int i = 1; i <= n; ++i) {
            for (int j = 1; j <= k; ++j) {
                // don't take here
                dpL1[i][j] = dpL1[i-1][j-1];
                dpL2[i][j] = dpL2[i-1][max(0, j-2)];

                // take here
                dpL1[i][j] = max(dpL1[i][j], a[i] + dpL1[i-1][max(0, j-2)]);
                dpL2[i][j] = max(dpL2[i][j], a[i] + dpL2[i-1][max(0, j-3)]);

                // allow <= j
                dpL1[i][j] = max(dpL1[i][j], dpL1[i][j-1]);
                dpL2[i][j] = max(dpL2[i][j], dpL2[i][j-1]);
            }
        }

        for (int i = n; i >= 1; --i) {
            for (int j = 1; j <= k; ++j) {
                // don't take here
                dpR1[i][j] = dpR1[i+1][j-1];
                dpR2[i][j] = dpR2[i+1][max(0, j-2)];

                // take here
                dpR1[i][j] = max(dpR1[i][j], a[i] + dpR1[i+1][max(0, j-2)]);
                dpR2[i][j] = max(dpR2[i][j], a[i] + dpR2[i+1][max(0, j-3)]);

                // allow <= j
                dpR1[i][j] = max(dpR1[i][j], dpR1[i][j-1]);
                dpR2[i][j] = max(dpR2[i][j], dpR2[i][j-1]);
            }
        }

        for (int i = 1; i <= n; ++i) {
            int ans = 0;
            for (int x = 0; x <= k; ++x) {
                ans = max(ans, dpL2[i][x] + dpR1[i+1][max(0, k-x-1)]);
                ans = max(ans, dpR2[i][x] + dpL1[i-1][max(0, k-x-1)]);
            }
            cout << ans << ' ';
        }
        cout << '\n';
    }
}
2

Do Java solutions pass?