Leetcode 0265. Paint House II

Problem statement

https://leetcode.com/problems/paint-house-ii/

Solution

Let dp[i][j] be the minimum cost to paint houses upto number i such that last house has color j. Then dp[i][j] = $\min\limits_{k\neq j}$ dp[i-1][k] + costs[i][j]. To be able to update this in $O(1)$ we need to keep two smallest cells from $i-1$-the column of our table, which can be done in $O(k)$ time.

Complexity

Time complexity is $O(nk)$, space complexity can be done in $O(k)$, if we use only last row of matrix, as I did in the following code.

Code

class Solution:
    def minCostII(self, costs):
        n, k = len(costs), len(costs[0])
        if k == 1 and n == 1: return costs[0][0]
        dp = [0]*k
        
        for cost in costs:
            dp_new = [0]*k
            min_ind = dp.index(min(dp))
            for i in range(k):
                if i != min_ind:
                    dp_new[i] = dp[min_ind] + cost[i]
                else:
                    dp_new[i] = min(dp[i] for i in range(k) if i != min_ind) + cost[i]
                    
            dp = dp_new
            
        return min(dp)