Leetcode 23. Merge k Sorted Lists with Python

문제

You are given an array of k linked-lists lists, 
each linked-list is sorted in ascending order.

Merge all the linked-lists into one sorted linked-list and return it.

k개의 링크드 리스트가 주어집니다.
이들을 합쳐 하나의 오름차순으로 정렬된 리스트를 만들어 리턴하세요.

예시

Input: lists = [[1,4,5],[1,3,4],[2,6]]
Output: [1,1,2,3,4,4,5,6]

제한

- k == lists.length
- 0 <= k <= 104
- 0 <= lists[i].length <= 500
- -10^4 <= lists[i][j] <= 10^4
- lists[i] 는 오름차순으로 정렬됩니다.
- The sum of lists[i].length will not exceed 104.

풀이법

값을 하나씩 불러와 정렬할수 있는것은 Heap 정렬이 있습니다.

Heap 정렬을 사용합니다.

from typing import List, Optional


class ListNode:
    def __init__(self, val=0, next=None):
        self.val = val
        self.next = next
        
class heap:
    def __init__(self):
        self.a = [0] * 5000001
        self.t = 0
    def put(self, num):
        self.t += 1
        c = self.t
        while c > 1:
            if self.a[c//2] <= num: break
            self.a[c], c = self.a[c//2], c//2
        self.a[c] = num
    def pull(self):
        if self.t == 0: return None
        ret = self.a[1]
        self.a[1], self.a[self.t] = self.a[self.t], self.a[1]
        self.t -= 1
        key = self.a[1]
        parent = 1
        child = 2
        while child <= self.t:
            if child + 1 <= self.t and self.a[child] > self.a[child + 1]:
                child += 1
            if self.a[child] > key: break
            self.a[parent] = self.a[child]
            parent = child
            child = parent * 2
        self.a[parent] = key
        return ret
    def __str__(self):
        return str(self.a)
    
class Solution:
    def mergeKLists(self, lists: List[Optional[ListNode]]) -> Optional[ListNode]:
        h = heap()
        retList = ListNode()
        head = retList
        for lst in lists:
            while lst != None:
                h.put(lst.val)
                lst = lst.next
        while True:
            d = h.pull()
            if d != None: 
                retList.next = ListNode(d, None)
                retList = retList.next
            else: return head.next