☕ DSA Java Articles

Topics › Dynamic Programming › 2D/3D DP and DP on Grids

Ninja's training

Medium TUF ID 292 A2Z

Java source ↗ Edit this article ↗

Problem

A ninja trains for n days. Each day he does exactly one of three activities (say running, fighting practice, learning new moves), and points[day][task] is the merit he earns for doing activity task on that day. He refuses to do the same activity on two consecutive days. Return the maximum total merit he can earn.

Constraints (Coding Ninjas version): 1 ≤ n ≤ 10^5, 1 ≤ points[day][task] ≤ 100. The total is at most 10^7, so int is fine. The Java methods also accept n = 0 and return 0.

Example 1

points = [[1, 2, 5],
          [3, 1, 1],
          [3, 3, 3]]
answer = 11         // day 0 task 2 (5), day 1 task 0 (3), day 2 task 1 or 2 (3)

Example 2

points = [[10, 40, 70],
          [20, 50, 80],
          [30, 60, 90]]
answer = 210        // 70 + 50 + 90

Why greedy fails

Picking the best allowed activity each day, starting from day 0, looks natural. On [[10, 50, 1], [5, 100, 11]] greedy takes 50 on day 0, which forbids the 100 on day 1, and ends with 50 + 11 = 61. Taking 10 on day 0 leaves the 100 free: 110. Today’s best choice depends on what tomorrow offers, so we must compare whole schedules.

The state

In the 1D problems the state was a single index. Here the index alone is not enough. Knowing that we are on day day does not tell us which activities are allowed, because that depends on the neighbouring day’s choice. So the state carries one more piece of information:

f(day, last) = the best total for days 0 .. day, given that activity last is forbidden on day (because it is what the ninja does on day day + 1).

last ranges over 0, 1, 2, plus a fourth value 3 meaning “nothing is forbidden”, used for the final day, which has no tomorrow. The answer is f(n - 1, 3).

The transition tries every allowed activity on day and passes it down as the forbidden one for day - 1:

f(-1, any)  = 0
f(day, last) = max over task ≠ last of  points[day][task] + f(day - 1, task)

Approach 1: Recursion

Intuition. Work from the last day backwards. On each day try the (two or three) activities that differ from the following day’s, add their points, and recurse into the previous day with that activity forbidden. The best of these is the answer for the state. Every valid schedule is reached by exactly one path of choices, and no invalid one can be, so the maximum is correct.

Algorithm.

  1. If day < 0, return 0.
  2. For each task from 0 to 2 except last, compute points[day][task] + f(day - 1, task).
  3. Return the maximum.
static int recursive(int[][] points) {
    return solve(points.length - 1, NONE, points);
}

static int solve(int day, int last, int[][] points) {
    if (day < 0) return 0;                                  // no days left
    int best = 0;
    for (int task = 0; task < 3; task++) {
        if (task == last) continue;                         // same task as the day after: not allowed
        best = Math.max(best, points[day][task] + solve(day - 1, task, points));
    }
    return best;
}

Complexity. Time O(2^n): the top call branches three ways and every later call two ways, so there are about 3 · 2^(n-1) leaves. Space O(n) for the recursion depth.

Approach 2: Memoization

Intuition. There are only n × 4 distinct states. The recursion keeps revisiting them: f(day - 2, 0) is requested by both f(day - 1, 1) and f(day - 1, 2), and each of those is requested by several states above it, so the repeats multiply. Store each answer in dp[day][last].

Algorithm.

  1. Create dp[n][4] filled with -1; real totals are never negative.
  2. In the helper, return the cached value when present; otherwise compute it as above and store it.
static int memoization(int[][] points) {
    int n = points.length;
    int[][] dp = new int[n][4];
    for (int[] row : dp) Arrays.fill(row, -1);              // -1 = not computed (totals are >= 0)
    return memo(n - 1, NONE, points, dp);
}

static int memo(int day, int last, int[][] points, int[][] dp) {
    if (day < 0) return 0;
    if (dp[day][last] != -1) return dp[day][last];
    int best = 0;
    for (int task = 0; task < 3; task++) {
        if (task == last) continue;
        best = Math.max(best, points[day][task] + memo(day - 1, task, points, dp));
    }
    return dp[day][last] = best;
}

Complexity. Time O(n · 4 · 3) = O(n): each of the 4n states loops over three tasks once. Space O(n) for the table plus O(n) recursion stack.

Approach 3: Tabulation

Intuition. f(day, ·) only reads f(day - 1, ·), so fill the table one day at a time, starting from day 0. For day 0 there is no earlier day, so dp[0][last] is just the best single activity other than last.

Algorithm.

  1. For each last in 0..3: dp[0][last] = max of points[0][task] for task ≠ last.
  2. For day from 1 to n - 1, for each last in 0..3: dp[day][last] = max over task ≠ last of points[day][task] + dp[day - 1][task].
  3. Return dp[n - 1][3].
static int tabulation(int[][] points) {
    int n = points.length;
    if (n == 0) return 0;
    int[][] dp = new int[n][4];
    for (int last = 0; last < 4; last++) {                  // day 0: best single task other than `last`
        for (int task = 0; task < 3; task++) {
            if (task != last) dp[0][last] = Math.max(dp[0][last], points[0][task]);
        }
    }
    for (int day = 1; day < n; day++) {
        for (int last = 0; last < 4; last++) {
            for (int task = 0; task < 3; task++) {
                if (task != last) dp[day][last] = Math.max(dp[day][last], points[day][task] + dp[day - 1][task]);
            }
        }
    }
    return dp[n - 1][NONE];
}

Dry run on Example 1. Column last = k means task k is not allowed on that day.

day points last = 0 last = 1 last = 2 last = 3
0 [1, 2, 5] 5 5 2 5
1 [3, 1, 1] 6 8 8 8
2 [3, 3, 3] 11 11 11 11

For example dp[1][1] = max(3 + dp[0][0], 1 + dp[0][2]) = max(3 + 5, 1 + 2) = 8: do task 0 on day 1 (so task 0 is forbidden on day 0, leaving the 5), or task 2 (leaving only 1 or 2 on day 0). The answer is dp[2][3] = 11.

Complexity. Time O(n · 4 · 3), space O(n · 4).

Approach 4: Space optimized

Intuition. Row day reads only row day - 1. Keep a single array prev[4]. Starting it at all zeros is the same as the base case f(-1, any) = 0, which means day 0 needs no special handling at all: its row is computed by the same loop as every other day.

Algorithm.

  1. prev = [0, 0, 0, 0].
  2. For each day’s row today: build cur[last] = max over task ≠ last of today[task] + prev[task], then set prev = cur.
  3. Return prev[3].
static int spaceOptimized(int[][] points) {
    int[] prev = new int[4];                                // prev[last] = best total for the days so far (0 before day 0)
    for (int[] today : points) {
        int[] cur = new int[4];
        for (int last = 0; last < 4; last++) {
            for (int task = 0; task < 3; task++) {
                if (task != last) cur[last] = Math.max(cur[last], today[task] + prev[task]);
            }
        }
        prev = cur;
    }
    return prev[NONE];
}

Complexity. Time O(n · 4 · 3), space O(1) (two arrays of four integers).

Edge cases and pitfalls