Chương trình tìm k danh sách con có tổng lớn nhất và trả về tổng theo thứ tự tăng dần trong Python

Mô tả bài toán

Cho danh sách nums và số k. Tìm k danh sách con (liên tiếp) có tổng lớn nhất, trả về các tổng theo thứ tự không giảm.

Ví dụ: nums = [2,4,5,-100,12,-30,6,-2,6], k=3 → đầu ra [10, 11, 12] (tương ứng sublist [6,-2,6], [2,4,5], [12]).

Giải thuật (Prefix Sum + Max Heap)

1. Tính prefix sum ps.
2. Đẩy tất cả tổng sublist -(ps[j]-ps[i]) vào max-heap (dùng heapq với dấu âm).
3. Lấy k phần tử lớn nhất từ heap, đổi dấu và đảo ngược để có thứ tự tăng dần.

Ví dụ mã Python

from heapq import heappop, heappush

class Solution:
    def solve(self, nums, k):
        ps = [0 for _ in range(len(nums) + 1)]
        for i, v in enumerate(nums):
            ps[i + 1] = v + ps[i]
        hp = []
        for i in range(len(ps)):
            for j in range(i + 1, len(ps)):
                heappush(hp, -(ps[j] - ps[i]))
        return list(reversed([-heappop(hp) for _ in range(k)]))

ob = Solution()
nums = [2, 4, 5, -100, 12, -30, 6, -2, 6]
k = 3
print(ob.solve(nums, k))

Đầu vào

[2, 4, 5, -100, 12, -30, 6, -2, 6],3

Đầu ra

[10, 11, 12]