Chương trình tìm chi phí nhỏ nhất sơn rào với k màu trong Python

Cho ma trận N x K chi phí sơn: hàng n là rào thứ n, cột k là màu k. Sơn sao cho 2 rào liền kề khác màu, chi phí tối thiểu.

Ví dụ ma trận 4x3 → kết quả 14 (chọn màu: 5→2→3→4).

Giải thuật (DP tối ưu O(N*K) theo dõi 2 giá trị nhỏ nhất)

Thay vì duyệt tất cả màu trước, chỉ lưu:

• fc, ft: màu và chi phí nhỏ nhất hàng trước.
• sc, st: màu và chi phí nhỏ nhì hàng trước.

Với hàng hiện tại, tính chi phí chọn màu i: cost = row[i] + (ft if i != fc else st). Cập nhật nfc, nft, nsc, nst cho hàng tiếp theo.

Ví dụ minh họa

class Solution:
    def solve(self, matrix):
        n = len(matrix)
        fc, ft = 0, 0       # first color, first cost
        sc, st = 1, 0       # second color, second cost
        inf = int(1e18)
        
        for row in matrix:
            nfc, nft = -1, inf
            nsc, nst = -1, inf
            for i, t in enumerate(row):
                ct = t + (ft if i != fc else st)
                if ct <= nft:
                    nsc, nst = nfc, nft
                    nfc, nft = i, ct
                elif ct <= nst:
                    nsc, nst = i, ct
            fc, ft = nfc, nft
            sc, st = nsc, nst
        return ft

ob = Solution()
matrix = [[6,4,5],[3,2,7],[3,4,5],[5,4,4]]
print(ob.solve(matrix))  # Output: 14

Đầu vào / Đầu ra

Input: [[6,4,5],[3,2,7],[3,4,5],[5,4,4]]

Output: 14