ML Prep

Merge sorted lists with a heap

Problem

Implement merge_sorted_lists(rows) where every inner list is sorted. Return one sorted list. Use a heap containing at most one pending value from each nonempty row, so NN total values across kk rows take O(Nlog⁡k)O(N \log k) time.

Reveal answer or reference solution

Seed the heap with each row's first value and its row and column indices. Repeatedly pop the smallest entry, append its value, and push the next value from the same row. Row and column indices break value ties deterministically.

import heapq
 
def merge_sorted_lists(rows: list[list[int]]) -> list[int]:
    heap: list[tuple[int, int, int]] = []
    for row_index, row in enumerate(rows):
        if row:
            heapq.heappush(heap, (row[0], row_index, 0))
    merged: list[int] = []
    while heap:
        value, row_index, column_index = heapq.heappop(heap)
        merged.append(value)
        next_column = column_index + 1
        if next_column < len(rows[row_index]):
            next_value = rows[row_index][next_column]
            heapq.heappush(
                heap, (next_value, row_index, next_column)
            )
    return merged

Public tests

  • merge_sorted_lists([]) → []
  • merge_sorted_lists([[], []]) → []
  • merge_sorted_lists([[1, 4], [1, 3, 5], [2]]) → [1, 1, 2, 3, 4, 5]
  • merge_sorted_lists([[-3, 7], [-2, 0]]) → [-3, -2, 0, 7]
Outcome

Previous attempts

Loading attempts saved in this browser…

    Review Screening Results