LeetCode - Paint House

These problems are basically same except the constraints. In the first question, k is always 3. If we solve the second problem, then we can also solve the first problem. Therefore, let's think about the case when k is not fixed.

Paint House (Medium)

Paint House (Medium)

There is a row of n houses, where each house can be painted one of three colors: red, blue, or green. The cost of painting each house with a certain color is different. You have to paint all the houses such that no two adjacent houses have the same color.

The cost of painting each house with a certain color is represented by an n x 3 cost matrix costs.

For example, costs[0][0] is the cost of painting house 0 with the color red; costs[1][2] is the cost of painting house 1 with color green, and so on...
Return the minimum cost to paint all houses.

Paint House II (Hard)

Paint House II (Hard)

There are a row of n houses, each house can be painted with one of the k colors. The cost of painting each house with a certain color is different. You have to paint all the houses such that no two adjacent houses have the same color.

The cost of painting each house with a certain color is represented by an n x k cost matrix costs.

For example, costs[0][0] is the cost of painting house 0 with color 0; costs[1][2] is the cost of painting house 1 with color 2, and so on...
Return the minimum cost to paint all houses.

Solution 1 : Dynamic Programming

Supposing we got 5 houses and 6 colors (red, green, blue, yellow, purple, orange).

Starting from the second, we can only choose the different color from the previous row. For example, if costs[1][0] is selected, then the only possible color to be chosen from the previous row must be from column 1 to column 5 and pick the minimum costs[0][j] where $j \neq 0$ in this example.

Therefore, we iterate each color in the current row and compare it from the previous row to find the minimum cost. The following solution gives $O(n * k ^ 2)$ time complexity and $O(k)$ space complexity.

def minCostII(self, costs: List[List[int]]) -> int:
    n = len(costs)
    k = len(costs[0])
    prev_row = costs[0]
    for house in range(1, n):
        cur_row = [0] * k
        for cur_color in range(k):
            prev_best = math.inf
            for prev_color in range(k):
                if cur_color == prev_color:
                    continue
                prev_best = min(prev_best, prev_row[prev_color])
            cur_row[cur_color] += prev_best + costs[house][cur_color]
        prev_row = cur_row
    return min(prev_row)

which can be written in a cleaner way.

def minCostII(self, costs: List[List[int]]) -> int:
    k = len(costs[0])
    prev_row = [0] * k
    for cur_row in costs:
        prev_row = [cur_row[i] + min(prev_row[:i] + prev_row[i + 1:]) for i in range(k)]
    return min(prev_row)

Solution 2: Markov Chain

If n and k are set to $500$, then Solution 1 will exceed time limit. Here's another solution using the idea of Markov Chain. Given that k states (colors) and n stages, we can update the matrix to find out the optimal value for each state at the current stage. The idea is to find out the first and the second minimum cost for each stage. For the next stage, we can update each cost by adding either the first minimum cost or the second one. The reason we need two minimum cost is that we cannot paint the same color. If the index of the first minimum cost from the previous stage is i and we have to use the second minimum cost for the current stage, and vice versa.

class Solution:
    def solve(self, matrix):
        n = len(matrix)
        k = len(matrix[0])
        for house in range(1, n):
            first_mi_color = second_mi_color = None
            # find first & second min color
            for color in range(k):
                cost = matrix[house - 1][color]
                if first_mi_color is None or cost < matrix[house - 1][first_mi_color]:
                    second_mi_color = first_mi_color
                    first_mi_color = color
                elif second_mi_color is None or cost < matrix[house - 1][second_mi_color]:
                    second_mi_color = color
            # update matrix for the current row
            for color in range(k):
                if color == first_mi_color:
                    matrix[house][color] += matrix[house - 1][second_mi_color]
                else:
                    matrix[house][color] += matrix[house - 1][first_mi_color]
        return min(matrix[-1])

Here's the shorter version.

class Solution:
    def minCostII(self, costs: List[List[int]]) -> int:
        n, k = len(costs), len(costs[0])
        for house in range(1, n):
            first_mi_cost = min(costs[house - 1])
            idx = costs[house - 1].index(first_mi_cost)
            second_mi_cost = min(costs[house - 1][:idx] + costs[house - 1][idx + 1:])
            for color in range(k):
                if color == idx:
                    costs[house][color] += second_mi_cost
                else:
                    costs[house][color] += first_mi_cost
        return min(costs[-1])