Jan 1, 2025

Dynamic Programming - 2D&More

Overview

“2D” describes the state, not necessarily the shape of the input. A subproblem needs two coordinates when one index is not enough to describe it. Those coordinates may be a grid cell, two prefix lengths, a house and its color, or a transaction count and holding status. Some problems need a third coordinate, such as the remainder of a path sum.

The main patterns in this post are:

Grid or cell DP: dp[r][c] describes a result at cell (r, c). Its predecessors may be above, left, or diagonally adjacent. This covers path counting, maximal squares, and falling paths.
Two-prefix DP: dp[i][j] describes the relationship between the first i elements of one sequence and the first j elements of another. Matching and edit-distance problems usually use the top, left, and diagonal states.
Position plus state: one coordinate is time or position and the other is a finite choice such as color, transaction count, or whether a stock position is open. Paint House and the stock problems follow this model.
Additional state: add another dimension only when different values cannot safely share one state. For example, paths with different sum remainders must remain separate until the destination.

Define every state in a complete sentence before writing its recurrence. Then identify base cases, valid predecessors, and an iteration order in which those predecessors have already been computed. Padding the table with an empty row or column often makes prefix and grid boundaries simpler. A full table may be compressed to the previous row or a few state arrays when future transitions no longer need older entries; use snapshots when all transitions for the current step must read only from the previous step.

The time and space complexity are usually the product of the state dimensions. An R × C × K state space therefore typically costs O(RCK). Also check whether DP is necessary at all: Minimum Window Subsequence in this post uses a forward match and backward contraction instead of storing every pair of prefix states.

Templates

Grid or cell DP

Define exactly what value ends at or reaches (r, c).
Initialize boundary cells, then traverse in dependency order.
Combine only predecessors that can legally transition into the current cell.

class GridDpTemplate {
    public int countGridPaths(int rows, int columns) {
        int[][] dp = new int[rows][columns];

        for (int r = 0; r < rows; r++) {
            for (int c = 0; c < columns; c++) {
                if (r == 0 || c == 0) {
                    dp[r][c] = 1;
                } else {
                    dp[r][c] = dp[r - 1][c] + dp[r][c - 1];
                }
            }
        }

        return dp[rows - 1][columns - 1];
    }
}

def count_grid_paths(rows: int, columns: int) -> int:
    dp = [[0] * columns for _ in range(rows)]

    for r in range(rows):
        for c in range(columns):
            if r == 0 or c == 0:
                dp[r][c] = 1
            else:
                dp[r][c] = dp[r-1][c] + dp[r][c-1]

    return dp[-1][-1]

Two-prefix DP

Let dp[i][j] describe the answer for first[:i] and second[:j].
Initialize the empty-prefix row and column.
A matching pair usually uses the diagonal; otherwise consider skipping or changing one side through the top, left, and diagonal states.

class PrefixDpTemplate {
    public int minEditDistance(String first, String second) {
        int[][] dp = new int[first.length() + 1][second.length() + 1];

        for (int i = 0; i <= first.length(); i++) {
            dp[i][0] = i;
        }
        for (int j = 0; j <= second.length(); j++) {
            dp[0][j] = j;
        }

        for (int i = 1; i <= first.length(); i++) {
            for (int j = 1; j <= second.length(); j++) {
                if (first.charAt(i - 1) == second.charAt(j - 1)) {
                    dp[i][j] = dp[i - 1][j - 1];
                } else {
                    dp[i][j] = 1 + Math.min(
                        dp[i - 1][j - 1],
                        Math.min(dp[i - 1][j], dp[i][j - 1])
                    );
                }
            }
        }

        return dp[first.length()][second.length()];
    }
}

def min_edit_distance(first: str, second: str) -> int:
    dp = [[0] * (len(second) + 1) for _ in range(len(first) + 1)]

    for i in range(len(first) + 1):
        dp[i][0] = i
    for j in range(len(second) + 1):
        dp[0][j] = j

    for i in range(1, len(first) + 1):
        for j in range(1, len(second) + 1):
            if first[i-1] == second[j-1]:
                dp[i][j] = dp[i-1][j-1]
            else:
                dp[i][j] = 1 + min(dp[i-1][j-1], dp[i-1][j], dp[i][j-1])

    return dp[-1][-1]

Position plus state

Keep one value for every meaningful state at the current position.
Derive the next values only from the previous position.
Use copies when updating in place could let one transition incorrectly feed another transition during the same step.

class StateDpTemplate {
    public int maxProfit(int k, int[] prices) {
        int[] buy = new int[k + 1];
        int[] sell = new int[k + 1];
        for (int t = 0; t <= k; t++) {
            buy[t] = Integer.MIN_VALUE / 2;
        }

        for (int price : prices) {
            int[] oldBuy = buy.clone();
            int[] oldSell = sell.clone();

            for (int t = 1; t <= k; t++) {
                buy[t] = Math.max(oldBuy[t], oldSell[t - 1] - price);
                sell[t] = Math.max(oldSell[t], oldBuy[t] + price);
            }
        }

        return sell[k];
    }
}

def max_profit(prices: list[int], k: int) -> int:
    buy = [float("-inf")] * (k + 1)
    sell = [0] * (k + 1)

    for price in prices:
        old_buy = buy.copy()
        old_sell = sell.copy()

        for t in range(1, k + 1):
            buy[t] = max(old_buy[t], old_sell[t-1] - price)
            sell[t] = max(old_sell[t], old_buy[t] + price)

    return sell[k]

Problems

62. Unique Paths

Count paths from the top-left to the bottom-right when movement is limited to right and down. Every cell can be reached from its top or left neighbor, and every first-row or first-column cell has exactly one path.

class Solution {
    public int uniquePaths(int m, int n) {
        int[][] grid = new int[m][n]; // m rows, n columns

        for (int r = 0; r < m; r++) {
            for (int c = 0; c < n; c++) {
                if (r == 0 || c == 0) {
                    grid[r][c] = 1;
                    continue;
                }

                grid[r][c] = grid[r - 1][c] + grid[r][c - 1];
            }
        }

        return grid[m - 1][n - 1];
    }
}

class Solution:
    def uniquePaths(self, m: int, n: int) -> int:
        grid = [[0]*n for _ in range(m)] # m rows, n columns

        for r in range(m):
            for c in range(n):
                if r == 0 or c == 0:
                    grid[r][c] = 1
                    continue

                grid[r][c] = grid[r-1][c] + grid[r][c-1]

        return grid[-1][-1]

10. Regular Expression Matching

Let dp[i][j] indicate whether the first i characters of s match the first j characters of p. A normal character or . uses the diagonal state. For *, either ignore the preceding pattern element or consume one matching character while keeping that element available.

class Solution {
    public boolean isMatch(String s, String p) {
        int Ls = s.length();
        int Lp = p.length();

        // dp[i][j] = whether s.substring(0, i) matches p.substring(0, j)
        boolean[][] dp = new boolean[1 + Ls][1 + Lp];
        dp[0][0] = true;

        // Empty s can match patterns like a*, a*b*, a*b*c*
        // Note that only p contains patterns like .*
        for (int j = 2; j <= Lp; j++) {
            if (p.charAt(j - 1) == '*') {
                dp[0][j] = dp[0][j - 2];
            }
        }

        for (int i = 1; i <= Ls; i++) {
            for (int j = 1; j <= Lp; j++) {
                if (p.charAt(j - 1) == '.' || p.charAt(j - 1) == s.charAt(i - 1)) {
                    dp[i][j] = dp[i - 1][j - 1];
                } else if (p.charAt(j - 1) == '*') {
                    // Case 1: use x* zero times, ignore this pattern completely
                    dp[i][j] = dp[i][j - 2];

                    // Case 2: use x* one or more times (s = "aa", p = "aa*")
                    if (p.charAt(j - 2) == '.' || p.charAt(j - 2) == s.charAt(i - 1)) {
                        dp[i][j] = dp[i][j] || dp[i - 1][j];
                    }
                }
            }
        }

        return dp[Ls][Lp];
    }
}

class Solution:
    def isMatch(self, s: str, p: str) -> bool:
        Ls, Lp = len(s), len(p)

        # dp[i][j] = whether s[:i] matches p[:j], similar to edit distance
        dp = [[False]*(1 + Lp) for _ in range(1 + Ls)]
        dp[0][0] = True

        # Empty s can match patterns like a*, a*b*, a*b*c*
        # Note that ONLY p contains the pattern like . *
        for j in range(2, 1 + Lp):
            if p[j - 1] == '*':
                dp[0][j] = dp[0][j - 2]

        for i in range(1, Ls + 1):
            for j in range(1, Lp + 1):
                if p[j-1] == '.' or p[j-1] == s[i-1]:
                    dp[i][j] = dp[i - 1][j - 1]

                elif p[j-1] == '*':
                    # Case 1: use x* zero times, ignore this pattern completely
                    dp[i][j] = dp[i][j-2]

                    # Case 2: use x* one or more times (s = "aa", p = "aa*")
                    if p[j-2] == '.' or p[j-2] == s[i-1]:
                        dp[i][j] = dp[i][j] or dp[i-1][j]

        return dp[Ls][Lp]

44. Wildcard Matching

Use dp[i][j] to match prefixes of the input and pattern. A letter or ? consumes one character from both prefixes. A * either matches nothing by moving left in the table or consumes another input character by moving up.

class Solution {
    public boolean isMatch(String s, String p) {
        int Ls = s.length();
        int Lp = p.length();
        // dp[i][j]: can s.substring(0, i) match p.substring(0, j)?
        boolean[][] dp = new boolean[1 + Ls][1 + Lp];
        dp[0][0] = true;

        for (int j = 1; j <= Lp; j++) {
            if (p.charAt(j - 1) == '*') {
                dp[0][j] = dp[0][j - 1];
            }
        }

        // Good example: s = "aa", p = "a*"
        for (int i = 1; i <= Ls; i++) {
            for (int j = 1; j <= Lp; j++) {
                if (s.charAt(i - 1) == p.charAt(j - 1) || p.charAt(j - 1) == '?') {
                    dp[i][j] = dp[i - 1][j - 1];
                } else if (p.charAt(j - 1) == '*') {
                    dp[i][j] = dp[i - 1][j] || dp[i][j - 1];
                }
            }
        }

        return dp[Ls][Lp];
    }
}

class Solution:
    def isMatch(self, s: str, p: str) -> bool:
        Ls, Lp = len(s), len(p)
        # dp[i][j]: can s[:i] match p[:j]?
        dp = [[False] * (1 + Lp) for _ in range(1 + Ls)]
        dp[0][0] = True

        for j in range(1, 1 + Lp):
            if p[j-1] == "*":
                dp[0][j] = dp[0][j-1]

        # Good example: s = "aa", p = "a*"
        for i in range(1, 1 + Ls):
            for j in range(1, 1 + Lp):
                if s[i-1] == p[j-1] or p[j-1] == "?":
                    dp[i][j] = dp[i-1][j-1]
                elif p[j-1] == "*":
                    dp[i][j] = dp[i-1][j] or dp[i][j-1]

        return dp[-1][-1]

72. Edit Distance

Find the minimum insertions, deletions, and replacements required to transform one word into another. dp[i][j] is the cost for the first i and j characters. Equal final characters reuse the diagonal; otherwise, take one plus the best neighboring operation.

class Solution {
    public int minDistance(String word1, String word2) {
        int L1 = word1.length();
        int L2 = word2.length();
        // dp[r][c] distance between word1.substring(0, c) and word2.substring(0, r)
        int[][] dp = new int[1 + L2][1 + L1];

        for (int c = 1; c <= L1; c++) {
            dp[0][c] = c;
        }
        for (int r = 1; r <= L2; r++) {
            dp[r][0] = r;
        }

        for (int r = 1; r <= L2; r++) {
            for (int c = 1; c <= L1; c++) {
                if (word1.charAt(c - 1) == word2.charAt(r - 1)) {
                    dp[r][c] = dp[r - 1][c - 1];
                } else {
                    dp[r][c] = 1 + Math.min(
                        dp[r][c - 1],
                        Math.min(dp[r - 1][c], dp[r - 1][c - 1])
                    );
                }
            }
        }

        return dp[L2][L1];
    }
}

class Solution:
    def minDistance(self, word1: str, word2: str) -> int:
        L1, L2 = len(word1), len(word2)
        # dp[r][c] distance between word1[:c] and word2[:r]
        dp = [[0] * (1 + L1) for _ in range(1 + L2)]

        for c in range(1, 1 + L1):
            dp[0][c] = c
        for r in range(1, 1 + L2):
            dp[r][0] = r

        for r in range(1, 1 + L2):
            for c in range(1, 1 + L1):
                if word1[c-1] == word2[r-1]:
                    dp[r][c] = dp[r-1][c-1]
                else:
                    dp[r][c] = 1 + min(dp[r][c-1], dp[r-1][c], dp[r-1][c-1])

        return dp[-1][-1]

727. Minimum Window Subsequence

Scan forward through s1 until all of s2 has been matched as a subsequence. Then scan backward from that ending position to find the smallest possible start for the same match. Record the window and restart just after its left boundary so later candidates are not skipped.

class Solution {
    public String minWindow(String s1, String s2) {
        int L1 = s1.length();
        int L2 = s2.length();
        int bestStart = -1;
        int bestLength = Integer.MAX_VALUE;
        int right = 0;

        while (right < L1) {
            // Step 1: move right forward until s2 is matched as a subsequence
            int j = 0;
            // Scans through s1 so right < L1
            while (right < L1) {
                if (s1.charAt(right) == s2.charAt(j)) {
                    j++;
                    if (j == L2) {
                        break;
                    }
                }
                right++;
            }

            // Could not match s2 anymore
            if (right == L1) {
                break;
            }

            // Now s1[?...right] contains s2 as subsequence.
            // Step 2: move left backward to shrink the window.
            j = L2 - 1;
            int left = right;
            // Tries to match all of s2 backwards so j >= 0
            while (j >= 0) {
                if (s1.charAt(left) == s2.charAt(j)) {
                    j--;
                }
                left--;
            }

            // After loop, left moved one step before the real start
            left++;
            int currLength = right - left + 1;

            if (currLength < bestLength) {
                bestLength = currLength;
                bestStart = left;
            }

            // Step 3: continue searching after this left boundary
            right = left + 1;
        }

        return bestStart == -1
            ? ""
            : s1.substring(bestStart, bestStart + bestLength);
    }
}

class Solution:
    def minWindow(self, s1: str, s2: str) -> str:
        L1, L2 = len(s1), len(s2)
        best_left = -1
        best_len = float("inf")
        right = 0

        while right < L1:
            # Step 1: move right forward until s2 is matched as a subsequence
            j = 0
            # Scans through s1 so right < L1
            while right < L1:
                if s1[right] == s2[j]:
                    j += 1
                    if j == L2:
                        break
                right += 1

            # Could not match s2 anymore
            if right == L1:
                break

            # Now s1[?...right] contains s2 as subsequence.
            # Step 2: move left backward to shrink the window.
            j = L2-1
            left = right
            # Tries to match all of s2 backwards so j >= 0
            while j >= 0:
                if s1[left] == s2[j]:
                    j -= 1
                left -= 1

            # After loop, left moved one step before the real start
            left += 1
            curr_len = right - left + 1

            if curr_len < best_len:
                best_len = curr_len
                best_left = left

            # Step 3: continue searching after this left boundary
            right = left + 1

        return "" if best_left == -1 else s1[best_left:best_left+best_len]

221. Maximal Square

Find the largest all-ones square. If a cell contains one, its square side is one plus the smallest square ending above, left, or diagonally above-left. Tracking the largest side yields the area.

class Solution {
    public int maximalSquare(char[][] matrix) {
        int R = matrix.length;
        int C = matrix[0].length;
        int[][] dp = new int[1 + R][1 + C];
        int maxLen = 0;

        for (int r = 1; r <= R; r++) {
            for (int c = 1; c <= C; c++) {
                if (matrix[r - 1][c - 1] == '1') {
                    dp[r][c] = 1 + Math.min(
                        dp[r - 1][c],
                        Math.min(dp[r - 1][c - 1], dp[r][c - 1])
                    );
                } else {
                    dp[r][c] = 0;
                }

                maxLen = Math.max(maxLen, dp[r][c]);
            }
        }

        return maxLen * maxLen;
    }
}

from typing import List

class Solution:
    def maximalSquare(self, matrix: List[List[str]]) -> int:
        R, C = len(matrix), len(matrix[0])
        dp = [[0]*(1+C) for _ in range(1+R)]
        max_len = 0

        for r in range(1, 1+R):
            for c in range(1, 1+C):
                if matrix[r-1][c-1] == '1':
                    dp[r][c] = 1 + min(dp[r-1][c], dp[r-1][c-1], dp[r][c-1])
                else:
                    dp[r][c] = 0

                max_len = max(max_len, dp[r][c])

        return max_len*max_len

931. Minimum Falling Path Sum

Choose one cell per row while moving straight down or diagonally. Add each cell to the smallest reachable value from the previous row. The matrix itself can serve as the DP table, giving O(n²) time and O(1) auxiliary space.

class Solution {
    public int minFallingPathSum(int[][] matrix) {
        // dp[r][c] = matrix[r][c] + min(dp[r - 1][c], dp[r - 1][c - 1], dp[r - 1][c + 1])
        // Can use matrix as the dp table.
        int N = matrix.length;

        for (int r = 1; r < N; r++) {
            for (int c = 0; c < N; c++) {
                if (c == 0) {
                    matrix[r][c] += Math.min(matrix[r - 1][c], matrix[r - 1][c + 1]);
                } else if (c == N - 1) {
                    matrix[r][c] += Math.min(matrix[r - 1][c], matrix[r - 1][c - 1]);
                } else {
                    matrix[r][c] += Math.min(
                        matrix[r - 1][c - 1],
                        Math.min(matrix[r - 1][c], matrix[r - 1][c + 1])
                    );
                }
            }
        }

        int res = matrix[N - 1][0];
        for (int value : matrix[N - 1]) {
            res = Math.min(res, value);
        }
        return res;
    }
}

from typing import List

class Solution:
    def minFallingPathSum(self, matrix: List[List[int]]) -> int:
        # dp[r][c] = matrix[r][c] + min(dp[r-1][c], dp[r-1][c-1], dp[r-1][c+1])
        # can use matrix as dp table
        N = len(matrix)

        for r in range(1, N):
            for c in range(N):
                if c == 0:
                    matrix[r][c] += min(matrix[r-1][c], matrix[r-1][c+1])
                elif c == N-1:
                    matrix[r][c] += min(matrix[r-1][c], matrix[r-1][c-1])
                else:
                    matrix[r][c] += min(matrix[r-1][c-1], matrix[r-1][c], matrix[r-1][c+1])

        return min(matrix[-1])

2435. Paths in Matrix Whose Sum Is Divisible by K

Count right-and-down paths whose sum is divisible by k. Let dp[r][c][rem] count paths reaching a cell with remainder rem. Subtract the current cell’s value from rem to find the predecessor remainder required from the top and left cells.

class Solution {
    public int numberOfPaths(int[][] grid, int k) {
        // dp[r][c][rem]: number of paths to [r, c] with pathSum % k == rem
        // prevRem = (rem - grid[r][c]) % k
        int R = grid.length;
        int C = grid[0].length;
        int[][][] dp = new int[R][C][k];
        dp[0][0][grid[0][0] % k] = 1;
        int MOD = 1_000_000_007;

        for (int r = 0; r < R; r++) {
            for (int c = 0; c < C; c++) {
                if (r == 0 && c == 0) {
                    continue;
                }
                for (int rem = 0; rem < k; rem++) {
                    int prevRem = (rem - grid[r][c] % k + k) % k;
                    long ways = 0;
                    if (r > 0) {
                        ways += dp[r - 1][c][prevRem];
                    }
                    if (c > 0) {
                        ways += dp[r][c - 1][prevRem];
                    }
                    dp[r][c][rem] = (int) (ways % MOD);
                }
            }
        }

        return dp[R - 1][C - 1][0];
    }
}

from typing import List

class Solution:
    def numberOfPaths(self, grid: List[List[int]], k: int) -> int:
        # dp[i][j][rem]: # of paths to [i, j] with path_sum % k == rem
        # dp[i][j][rem] = dp[i-1][j][prev_rem] + dp[i][j-1][prev_rem]
        # prev_rem = (rem - grid[i][j]) % k
        R, C = len(grid), len(grid[0])
        dp = [[[0]*k for _ in range(C)] for _ in range(R)]
        dp[0][0][grid[0][0]%k] = 1
        MOD = 10**9 + 7

        for r in range(R):
            for c in range(C):
                if r == 0 and c == 0:
                    continue
                for rem in range(k):
                    prev_rem = (rem - grid[r][c]) % k
                    if r > 0:
                        dp[r][c][rem] += dp[r-1][c][prev_rem]
                    if c > 0:
                        dp[r][c][rem] += dp[r][c-1][prev_rem]
                    dp[r][c][rem] = (dp[r][c][rem] % MOD)

        return dp[-1][-1][0]

256. Paint House

Let dp[house][color] be the minimum cost after painting through house with that house assigned color. The preceding house must use one of the other two colors, so each state adds its current cost to the smaller compatible previous state.

class Solution {
    public int minCost(int[][] costs) {
        int N = costs.length;
        int[][] dp = new int[N][3];
        dp[0] = costs[0].clone();

        for (int i = 1; i < N; i++) {
            dp[i][0] = costs[i][0] + Math.min(dp[i - 1][1], dp[i - 1][2]);
            dp[i][1] = costs[i][1] + Math.min(dp[i - 1][0], dp[i - 1][2]);
            dp[i][2] = costs[i][2] + Math.min(dp[i - 1][1], dp[i - 1][0]);
        }

        return Math.min(dp[N - 1][0], Math.min(dp[N - 1][1], dp[N - 1][2]));
    }
}

from typing import List

class Solution:
    def minCost(self, costs: List[List[int]]) -> int:
        N = len(costs)
        dp = [[0]*3 for _ in range(N)]
        dp[0] = costs[0]

        for i in range(1, N):
            dp[i][0] = costs[i][0] + min(dp[i-1][1], dp[i-1][2])
            dp[i][1] = costs[i][1] + min(dp[i-1][0], dp[i-1][2])
            dp[i][2] = costs[i][2] + min(dp[i-1][1], dp[i-1][0])

        return min(dp[-1])

265. Paint House II

Generalize the state to k colors. Maintain a size-two heap containing the smallest and second-smallest total costs from the previous row. A color uses the smallest previous cost unless it came from the same color, in which case it uses the second smallest. The runtime is O(nk).

import java.util.PriorityQueue;

class Solution {
    public int minCostII(int[][] costs) {
        int N = costs.length;
        PriorityQueue<int[]> maxHeap = new PriorityQueue<>(
            (a, b) -> Integer.compare(a[0], b[0])
        );

        for (int idx = 0; idx < costs[0].length; idx++) {
            int[] pair = new int[]{-costs[0][idx], idx};
            if (idx < 2) {
                maxHeap.offer(pair);
            } else {
                maxHeap.offer(pair);
                maxHeap.poll();
            }
        }
        // Now maxHeap stores the lowest two costs for row 0.

        for (int r = 1; r < N; r++) {
            int[] pair1 = maxHeap.poll();
            int[] pair2 = maxHeap.poll();
            int cost1 = pair1[0];
            int idx1 = pair1[1];
            int cost2 = pair2[0];
            int idx2 = pair2[1];

            for (int idx = 0; idx < costs[r].length; idx++) {
                int val = costs[r][idx];
                // Pay attention to which cost is the smallest when using maxHeap.
                int newCost = idx != idx2 ? val - cost2 : val - cost1;
                int[] pair = new int[]{-newCost, idx};
                if (idx < 2) {
                    maxHeap.offer(pair);
                } else {
                    maxHeap.offer(pair);
                    maxHeap.poll();
                }
            }
        }

        int bestNegative = Integer.MIN_VALUE;
        for (int[] pair : maxHeap) {
            bestNegative = Math.max(bestNegative, pair[0]);
        }
        return -bestNegative;
    }
}

import heapq
from typing import List

class Solution:
    def minCostII(self, costs: List[List[int]]) -> int:
        N = len(costs)
        max_heap = []

        for idx, val in enumerate(costs[0]):
            if idx < 2:
                heapq.heappush(max_heap, (-val, idx))
            else:
                heapq.heappushpop(max_heap, (-val, idx))
        # now max_heap stores the lowest two costs for row 0

        for r in range(1, N):
            cost1, idx1 = heapq.heappop(max_heap)
            cost2, idx2 = heapq.heappop(max_heap)
            for idx, val in enumerate(costs[r]):
                # pay attention to which cost is the smallst when using max_heap
                new_cost = val - cost2 if idx != idx2 else val - cost1
                if idx < 2:
                    heapq.heappush(max_heap, (-new_cost, idx))
                else:
                    heapq.heappushpop(max_heap, (-new_cost, idx))

        return -max_heap[-1][0]

123. Best Time to Buy and Sell Stock III

Track the best profit after the first buy, first sale, second buy, and second sale. Each state either keeps its previous value or performs its corresponding action at today’s price. The four named states give O(n) time and O(1) space.

class Solution {
    public int maxProfit(int[] prices) {
        // All variables store the maximum profit for their state.
        int buy1 = Integer.MIN_VALUE;
        int sell1 = 0;
        int buy2 = Integer.MIN_VALUE;
        int sell2 = 0;

        for (int p : prices) {
            int oldBuy1 = buy1;
            buy1 = Math.max(buy1, -p);
            sell1 = Math.max(sell1, oldBuy1 + p);
            int oldBuy2 = buy2;
            buy2 = Math.max(buy2, sell1 - p);
            sell2 = Math.max(sell2, oldBuy2 + p);
        }

        return Math.max(sell1, sell2);
    }
}

from typing import List

class Solution:
    def maxProfit(self, prices: List[int]) -> int:
        # when I am at prices[i], depending on my status I can do different things
        # buy: could buy for -p sell: could sell for p
        # all the variables are the maximum profit for this state
        buy1, sell1 = float('-inf'), 0
        buy2, sell2 = float('-inf'), 0

        for p in prices:
            old_buy1 = buy1
            buy1 = max(buy1, -p)
            sell1 = max(sell1, old_buy1 + p)
            old_buy2 = buy2
            buy2 = max(buy2, sell1-p)
            sell2 = max(sell2, old_buy2 + p)

        return max(sell1, sell2)

188. Best Time to Buy and Sell Stock IV

Generalize the same holding and sold states to at most k transactions. Every day reads from snapshots of the previous day, so opening or closing a position cannot accidentally feed another transition on the same day. The runtime is O(nk) and the space is O(k).

class Solution {
    public int maxProfit(int k, int[] prices) {
        int[] buys = new int[1 + k];
        int[] sells = new int[1 + k];
        for (int j = 0; j <= k; j++) {
            buys[j] = Integer.MIN_VALUE / 2;
        }

        for (int p : prices) {
            int[] oldSells = sells.clone();
            int[] oldBuys = buys.clone();
            for (int j = 1; j <= k; j++) {
                buys[j] = Math.max(oldBuys[j], oldSells[j - 1] - p);
                sells[j] = Math.max(oldSells[j], oldBuys[j] + p);
            }
        }

        return sells[k];
    }
}

from typing import List

class Solution:
    def maxProfit(self, k: int, prices: List[int]) -> int:
        buys = [float('-inf')] * (1 + k)
        sells = [0] * (1 + k)

        for p in prices:
            old_sells = sells.copy()
            old_buys = buys.copy()
            for j in range(1, 1+k):
                buys[j] = max(old_buys[j], old_sells[j-1] - p)
                sells[j] = max(old_sells[j], old_buys[j] + p)

        return sells[-1]

3573. Best Time to Buy and Sell Stock V

For transaction t, track the best profit after closing it in sell[t], while holding a long position in buy[t], or while holding a short position in short[t]. Every transition reads from a snapshot of the previous day, preventing a position from being closed and reopened on the same day. Only closed states are valid final answers.

class Solution {
    public long maximumProfit(int[] prices, int k) {
        // Optional corner case handling
        if (prices.length <= 1 || k == 0) {
            return 0;
        }

        int K = Math.min(k, prices.length / 2);
        long negative = Long.MIN_VALUE / 4;
        long[] sell = new long[K + 1];
        long[] buy = new long[K + 1];
        long[] shortPosition = new long[k + 1];

        for (int t = 0; t <= K; t++) {
            buy[t] = negative;
        }
        for (int t = 0; t <= k; t++) {
            shortPosition[t] = negative;
        }

        for (int p : prices) {
            long[] oldSell = sell.clone();
            long[] oldBuy = buy.clone();
            long[] oldShort = shortPosition.clone();

            for (int t = 1; t <= K; t++) {
                buy[t] = Math.max(oldBuy[t], oldSell[t - 1] - p);
                shortPosition[t] = Math.max(oldShort[t], oldSell[t - 1] + p);
                sell[t] = Math.max(
                    oldSell[t],
                    Math.max(oldBuy[t] + p, oldShort[t] - p)
                );
            }
        }

        return sell[K]; // Do not return max(sell[K], shortPosition[K]).
    }
}

from typing import List

class Solution:
    def maximumProfit(self, prices: List[int], k: int) -> int:
        # Optional corner case handling
        if len(prices) <= 1 or k == 0:
            return 0

        K = min(k, len(prices) // 2)
        sell = [0] * (K + 1)
        buy = [float("-inf")] * (K + 1)
        short = [float("-inf")] * (k + 1)

        for p in prices:
            old_sell = sell.copy()
            old_buy = buy.copy()
            old_short = short.copy()

            for t in range(1, 1 + K):
                buy[t] = max(old_buy[t], old_sell[t-1] - p)
                short[t] = max(old_short[t], old_sell[t-1] + p)
                sell[t] = max(old_sell[t], old_buy[t] + p, old_short[t] - p)
        return sell[-1] # NOT max(sell[-1], short[-1])