Merge k sorted arrays or lists efficiently using a min-heap in O(n log k) time.
Published March 27, 2025
K-way merge solves problems involving multiple sorted data sources that need to be combined. A min-heap of size k efficiently picks the next smallest element in O(log k) time.
public ListNode mergeKLists(ListNode[] lists) {
// Min-heap: smallest node value at top
PriorityQueue<ListNode> heap = new PriorityQueue<>(
(a, b) -> a.val - b.val);
// Initialize: add the head of each non-empty list
for (ListNode node : lists)
if (node != null) heap.offer(node);
ListNode dummy = new ListNode(0), curr = dummy;
while (!heap.isEmpty()) {
ListNode node = heap.poll(); // smallest across all lists
curr.next = node;
curr = curr.next;
if (node.next != null) heap.offer(node.next); // add next from same list
}
return dummy.next;
}
// Time: O(n log k) where n = total nodes, k = number of lists
// Space: O(k) for the heap
public int[] mergeKSortedArrays(int[][] arrays) {
// Heap stores: [value, arrayIndex, elementIndex]
PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
int total = 0;
for (int i = 0; i < arrays.length; i++) {
if (arrays[i].length > 0) {
heap.offer(new int[]{arrays[i][0], i, 0});
total += arrays[i].length;
}
}
int[] result = new int[total];
int idx = 0;
while (!heap.isEmpty()) {
int[] top = heap.poll();
result[idx++] = top[0];
int arrIdx = top[1], elemIdx = top[2];
if (elemIdx + 1 < arrays[arrIdx].length)
heap.offer(new int[]{arrays[arrIdx][elemIdx+1], arrIdx, elemIdx+1});
}
return result;
}
public int kthSmallest(int[][] lists, int k) {
PriorityQueue<int[]> heap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
for (int i = 0; i < lists.length; i++)
if (lists[i].length > 0)
heap.offer(new int[]{lists[i][0], i, 0});
int count = 0;
while (!heap.isEmpty()) {
int[] top = heap.poll();
if (++count == k) return top[0];
int li = top[1], ei = top[2];
if (ei + 1 < lists[li].length)
heap.offer(new int[]{lists[li][ei+1], li, ei+1});
}
return -1;
}
// Find the smallest range [a, b] such that at least one number from each list is in [a, b]
public int[] smallestRange(List<List<Integer>> nums) {
PriorityQueue<int[]> minHeap = new PriorityQueue<>((a, b) -> a[0] - b[0]);
int max = Integer.MIN_VALUE;
for (int i = 0; i < nums.size(); i++) {
minHeap.offer(new int[]{nums.get(i).get(0), i, 0});
max = Math.max(max, nums.get(i).get(0));
}
int rangeStart = 0, rangeEnd = Integer.MAX_VALUE;
while (minHeap.size() == nums.size()) {
int[] curr = minHeap.poll();
if (max - curr[0] < rangeEnd - rangeStart) {
rangeStart = curr[0];
rangeEnd = max;
}
int li = curr[1], ei = curr[2];
if (ei + 1 < nums.get(li).size()) {
int next = nums.get(li).get(ei + 1);
minHeap.offer(new int[]{next, li, ei + 1});
max = Math.max(max, next);
}
}
return new int[]{rangeStart, rangeEnd};
}
[value, sourceIndex, positionInSource] — memorize this template.