Merge k Sorted Lists

Problem: Merge k Sorted Lists

We can use a heap to maintain heads of linked lists. Every time we extract a minimum, we replace the element in the heap by its next node in the linked list and heapify. When the heap is empty, the sorted lists will be merged.

Code in Python:

from heapq import *

# Definition for singly-linked list.
# class ListNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.next = None

class Solution(object):
    def mergeKLists(self, lists):
        """
        :type lists: List[ListNode]
        :rtype: ListNode
        """
        heap = []
        for head in lists:
            if head: heappush(heap, (head.val, head))
        prev = dummy = ListNode(0)
        while heap:
            val, node = heappop(heap)
            if node.next: heappush(heap, (node.next.val, node.next))
            prev.next, prev = node, node
        return dummy.next

In Java, we can use the class Priority Queue.

Code in Java:

/**
 * Definition for singly-linked list.
 * public class ListNode {
 *     int val;
 *     ListNode next;
 *     ListNode(int x) { val = x; }
 * }
 */
public class Solution {
    public ListNode mergeKLists(ListNode[] lists) {
        ListNode res = new ListNode(0);
        ListNode node = res;
        PriorityQueue<ListNode> pq = new PriorityQueue<ListNode>(new Comparator<ListNode>() {
            public int compare(ListNode n1, ListNode n2) {
                return n1.val-n2.val;
            }
        });

        for (ListNode head : lists) {
            if (head != null)
                pq.offer(head);
        }

        while (!pq.isEmpty()) {
            ListNode n = pq.poll();
            node.next = n;
            node = node.next;
            n = n.next;
            if (n != null) {
                pq.offer(n);
            }
        }

        node.next = null;
        return res.next;
    }
}

We can also use merge sort as divide and conquer to solve the problem, just like the solution here.