Chương trình tìm tổng lớn nhất K phần tử lấy từ các stack trong Python

Cho danh sách các stack và số . Lấy đúng phần tử (pop từ đỉnh các stack) để tổng lớn nhất.

Ví dụ: , → kết quả 39 (lấy 3 phần tử stack 1 và 1 phần tử stack 3: -15 + -4 + 50 + 8 = 39).

Giải thuật (Đệ quy + Memoization / DP)

Hàm : tổng max lấy từ stack trở đi với đã lấy phần tử.

• Base case: → 0; k hoặc → -inf.
• Nếu stack cuối: kiểm tra có đủ phần tử không, trả về tổng phần tử đỉnh.
• Thử lấy 0 đến phần tử từ stack , đệ quy stack tiếp theo.

Ví dụ minh họa

import math

class Solution:
    def solve(self, stacks, k):
        def rec(i, n):
            if n == k:
                return 0
            if n > k:
                return -math.inf
            if i == len(stacks):
                return -math.inf
            if i == len(stacks) - 1:
                needed = k - n
                if needed > len(stacks[i]):
                    return -math.inf
                else:
                    return sum(stacks[i][-needed:])
            res, su = -math.inf, 0
            for sti in range(len(stacks[i]) - 1, -1, -1):
                su += stacks[i][sti]
                localres = su + rec(i + 1, n + len(stacks[i]) - sti)
                res = max(res, localres)
            return max(res, rec(i + 1, n))
        return rec(0, 0)

ob = Solution()
stacks = [[50, -4, -15], [2], [6, 7, 8]]
k = 4
print(ob.solve(stacks, k))  # Output: 39

Đầu vào

[[50, -4, -15], [2], [6, 7, 8]], 4

Đầu ra

39