The Merge Champion

A medium Python interview practice problem on DataDriven. Write the Python and run it against test cases, with instant feedback.

Domain
Python
Difficulty
medium
Seniority
L4
Asked in
Citadel

Problem

We need every value from lists combined into one ascending list, where each inner list is a shard export that is already sorted on its own.

Many sorted rivers flowing into one.

Example 1

Input:

lists = [[1,4,5],[1,3,4],[2,6]]

Output:

[1,1,2,3,4,4,5,6]

Example 2

Input:

lists = [[]]

Output:

[]

Worked solution and explanation

What this really is

This is a merge of already-ordered streams wearing the costume of a list utility. Every shard hands you its order for free; the whole question is whether you spend it or throw it away. Anyone can produce an ascending list. What separates candidates is the shape of the merging: fold the shards in one by one and the early values get re-copied once per remaining shard, so the cost drifts toward n times k. Pair the shards up and merge in rounds, and every value is touched only log k times.

Building it

Step 1: Write the two-way merge as its own helper

merge_two walks left and right with one index each, unpacking the two current heads as a pair so the comparison reads as what it is. Using <= when the heads tie keeps the left value first. That changes nothing for plain numbers, but it keeps the merge stable if the values ever carry payloads.

Step 2: Merge in rounds, not left to right

rounds maps a round number to that round's lists, keyed by slot. Each round pads an odd count with one [] so every shard has a partner, then pairs the even slots with the odd slots as (left, right) tuples. A lone shard merges with nothing and passes through unchanged. Every round halves the count, and the finished round is popped so only one round's lists stay alive. After log k rounds, rounds[level][0] is the answer.

Step 3: Settle the empty cases up front

No shards at all returns [] immediately. Empty shards need no special handling: merging with an empty list just returns the other side, which is why [[]] and [[], [2, 4], [], [1, 3]] fall out correctly.

Pairwise merging in rounds
from typing import List


def merge_two(left: List[int], right: List[int]) -> List[int]:
    merged = []
    i, j = 0, 0
    while i < len(left) and j < len(right):
        left_head, right_head = left[i], right[j]
        if left_head <= right_head:
            merged.append(left_head)
            i += 1
        else:
            merged.append(right_head)
            j += 1
    merged.extend(left[i:])
    merged.extend(right[j:])
    return merged


def merge_k_sorted(lists: list) -> list:
    if not lists:
        return []

    # round number -> {slot: merged list}; only the live round is kept
    rounds = {0: dict(enumerate(lists))}
    level = 0
    while len(rounds[level]) > 1:
        current = list(rounds.pop(level).values())
        if len(current) % 2 == 1:
            current.append([])  # a lone shard merges with nothing
        pairs = list(zip(current[0::2], current[1::2]))
        next_round = {}
        for slot, (left, right) in enumerate(pairs):
            next_round[slot] = merge_two(left, right)
        level += 1
        rounds[level] = next_round

    return rounds[level][0]
Fold left to right
result = merge_two(result, shard) for each shard. The first shard's values are copied k times, the second's k minus 1 times, and so on: about n times k work.
Merge in rounds
Pair neighbours, merge, repeat. Each value moves once per round and there are log k rounds: n log k work with nothing but two indexes and a list.

Common follow-up questions