merge-k-sorted-lists.py

# Definition for singly-linked list.
# class ListNode:
#     def __init__(self, val=0, next=None):
#         self.val = val
#         self.next = next
class Solution:
    
    
    def mergeKLists(self, lists: List[ListNode]) -> ListNode:
        self.lists = lists
        
        if len(lists) == 0:
            return None

        return self._merge(0, len(lists) - 1)        
        
    def _merge(self, begin, end) -> ListNode:
       
        if begin == end:
            return self.lists[begin]
        
        if (end - begin) == 1:
            return self.combine(self.lists[begin], self.lists[end])
        
        middle = begin + (end - begin) // 2
        
        return self.combine(
            self._merge(begin, middle),
            self._merge(middle + 1, end)
        )
    
    def combine(self, A: ListNode, B: ListNode) -> ListNode:
        
        node = ListNode()
        current_node = node
        while A is not None and B is not None:
            
            if B.val < A.val:
                new_node = ListNode(val=B.val)
                B = B.next
            else:
                new_node = ListNode(val=A.val)
                A = A.next
            
            current_node.next = new_node
            current_node = new_node
            
        if A is not None:
            current_node.next = A
        if B is not None:
            current_node.next = B
        
        return node.next