Hợp nhất k danh sách đã sắp xếp trong Python

Cho k danh sách liên kết đã sắp xếp. Hợp nhất thành một danh sách đã sắp xếp. Sử dụng cấu trúc dữ liệu heap (priority queue). Ví dụ: [1,4,5], [1,3,4], [2,6] -> [1,1,2,3,4,4,5,6].

Các bước giải quyết:

• Tạo min-heap, đẩy nút đầu của mỗi danh sách vào heap
• Lặp: pop nút nhỏ nhất, thêm vào kết quả, đẩy nút tiếp theo của danh sách đó vào heap
• Lặp cho đến khi heap rỗng

Ví dụ

class ListNode:
    def __init__(self, data, next=None):
        self.val = data
        self.next = next

def make_list(elements):
    head = ListNode(elements[0])
    for element in elements[1:]:
        ptr = head
        while ptr.next:
            ptr = ptr.next
        ptr.next = ListNode(element)
    return head

def print_list(head):
    ptr = head
    print('[', end="")
    while ptr:
        print(ptr.val, end=", ")
        ptr = ptr.next
    print(']')

class Heap:
    def __init__(self):
        self.arr = []
    def insert(self, value):
        self.arr.append(value)
        n = len(self.arr) - 1
        i = n
        while i != 0 and self.arr[i].val < self.arr[self.parent(i)].val:
            self.arr[i], self.arr[self.parent(i)] = self.arr[self.parent(i)], self.arr[i]
            i = self.parent(i)
    def parent(self, i): return (i-1)//2
    def heapify(self, i):
        left = 2*i + 1
        right = 2*i + 2
        smallest = i
        n = len(self.arr)
        if left < n and self.arr[left].val < self.arr[smallest].val: smallest = left
        if right < n and self.arr[right].val < self.arr[smallest].val: smallest = right
        if smallest != i:
            self.arr[i], self.arr[smallest] = self.arr[smallest], self.arr[i]
            self.heapify(smallest)
    def extractMin(self):
        n = len(self.arr)
        if n == 0: return '#'
        if n == 1:
            temp = self.arr[0]
            self.arr.pop()
            return temp
        root = self.arr[0]
        self.arr[0] = self.arr[-1]
        self.arr.pop()
        self.heapify(0)
        return root

class Solution(object):
    def mergeKLists(self, lists):
        heap = Heap()
        for i in lists:
            if i:
                heap.insert(i)
        res = None
        res_next = None
        while True:
            temp = heap.extractMin()
            if temp == "#":
                return res
            if not res:
                res = temp
                res_next = temp
                temp = temp.next
                if temp:
                    heap.insert(temp)
                res.next = None
            else:
                res_next.next = temp
                temp = temp.next
                res_next = res_next.next
                if temp:
                    heap.insert(temp)
                res_next.next = None

ob = Solution()
lists = [[1,4,5],[1,3,4],[2,6]]
lls = []
for ll in lists:
    l = make_list(ll)
    lls.append(l)
print_list(ob.mergeKLists(lls))

Đầu vào

[[1,4,5],[1,3,4],[2,6]]

Đầu ra

[1, 1, 2, 3, 4, 4, 5, 6, ]