Sort an array with k worker threads

Question: Sort an array with k worker threads

Companies: xAI

Difficulty: Medium

Problem statement

Implement parallel_sort(nums: list[int], k: int) -> list[int], which returns every element of nums in ascending order and uses k worker threads for the heavy lifting.

Rules

• Split nums into contiguous segments, one per worker. Segment sizes may differ by at most one element.
• Workers run concurrently. A worker may read and write only the indices of its own segment.
• After all workers have finished, combine their results into the final sorted array. The combine step should take advantage of the work the workers already did. Don't simply sort everything again.
• Do not modify the caller's input list.

Constraints

• 0 <= len(nums) <= 10^6, 1 <= k <= 64, and values fit in a 32-bit signed integer.
• k may be larger than len(nums). Never start a thread for an empty segment.
• Duplicates and negative values are allowed.

Example 1

nums = [12, -4, 7, 7, 0, 30, 5], k = 3
segments: [12, -4, 7] | [7, 0] | [30, 5]
after workers: [-4, 7, 12] | [0, 7] | [5, 30]
result: [-4, 0, 5, 7, 7, 12, 30]

Example 2

nums = [8, 2], k = 6
only 2 workers are used: [8] | [2]
result: [2, 8]

State the time and space complexity of the worker phase and of the combine phase, and explain what makes the result safe to read once the workers are done.

Approach

The problem has two phases with a hard boundary between them: a parallel phase, where each worker may touch only its own slice, and a combine phase, which runs after every worker has finished.

Partition. Split n elements into k contiguous segments whose sizes differ by at most one. With q, r = divmod(n, k), the first r segments get q + 1 elements and the rest get q. If k > n, use only n workers, because an empty segment needs no thread. Segment i starts at i * q + min(i, r), so every worker can compute its own bounds without coordinating with the others.

Parallel phase. Each worker sorts nums[lo:hi] in place and writes only inside that range. The ranges don't overlap, so no locks are needed. The only synchronization is a join or barrier: the main thread waits until every worker is done before it reads the result. That "happens-before" edge is what makes the workers' writes visible, and it is the invariant to state out loud.

Combine phase. There are now k sorted runs. Some candidates concatenate them and sort again. That is correct but throws away the work the workers did and costs O(n log n) again. Merging the runs pairwise one after another, left to right, costs O(n·k). The intended approach is a k-way merge with a min-heap of (value, run_id, index): pop the smallest head, append it to the output, and push the next element from the same run. The heap never holds more than k entries, so the merge costs O(n log k). Ties break on run_id, which keeps the result stable across runs.

Invariants

• Every index in [0, n) belongs to exactly one segment, and the segments cover the array with no gaps.
• When the merge starts, every segment is sorted.
• The heap holds at most one cursor per run, and the popped values never decrease.

A note on Python: CPython's GIL means threads won't speed up a CPU-bound sort. Say so, then explain that the same structure runs on a ProcessPoolExecutor or in a language with real threads. Interviewers care about the structure, not the speedup.

Code

from __future__ import annotations

import heapq
import threading
from typing import List, Tuple


def partition(n: int, k: int) -> List[Tuple[int, int]]:
    """Split [0, n) into min(k, n) contiguous [lo, hi) ranges whose sizes differ by at most 1."""
    if k <= 0:
        raise ValueError("k must be positive")
    workers = min(k, n)
    if workers == 0:
        return []
    q, r = divmod(n, workers)
    bounds = []
    for i in range(workers):
        lo = i * q + min(i, r)
        hi = lo + q + (1 if i < r else 0)
        bounds.append((lo, hi))
    return bounds


def _sort_segment(nums: List[int], lo: int, hi: int) -> None:
    # Reads and writes only nums[lo:hi]; no other worker touches this range.
    nums[lo:hi] = sorted(nums[lo:hi])


def k_way_merge(nums: List[int], bounds: List[Tuple[int, int]]) -> List[int]:
    """Merge sorted runs nums[lo:hi] with a min-heap of run heads."""
    heap: List[Tuple[int, int, int]] = []  # (value, run_id, index)
    for run_id, (lo, hi) in enumerate(bounds):
        if lo < hi:
            heap.append((nums[lo], run_id, lo))
    heapq.heapify(heap)

    out: List[int] = []
    while heap:
        value, run_id, idx = heapq.heappop(heap)
        out.append(value)
        nxt = idx + 1
        if nxt < bounds[run_id][1]:
            heapq.heappush(heap, (nums[nxt], run_id, nxt))
    return out


def parallel_sort(nums: List[int], k: int) -> List[int]:
    """Sort nums with k workers, each confined to its own segment, then k-way merge."""
    work = list(nums)  # do not mutate the caller's list
    bounds = partition(len(work), k)

    threads = [
        threading.Thread(target=_sort_segment, args=(work, lo, hi))
        for lo, hi in bounds
    ]
    for t in threads:
        t.start()
    for t in threads:
        t.join()  # barrier: all segment writes are visible after this

    return k_way_merge(work, bounds)


if __name__ == "__main__":
    import random

    assert partition(10, 3) == [(0, 4), (4, 7), (7, 10)]
    assert partition(2, 5) == [(0, 1), (1, 2)]
    assert partition(0, 4) == []
    assert parallel_sort([5, 1, 4, 2, 3], 2) == [1, 2, 3, 4, 5]
    assert parallel_sort([], 3) == []
    assert parallel_sort([7], 4) == [7]
    assert parallel_sort([3, 3, 1, 1, 2, 2], 3) == [1, 1, 2, 2, 3, 3]
    assert parallel_sort([-2, 9, 0, -2, 5], 1) == [-2, -2, 0, 5, 9]
    for _ in range(200):
        arr = [random.randint(-50, 50) for _ in range(random.randint(0, 60))]
        assert parallel_sort(arr, random.randint(1, 12)) == sorted(arr)
    print("ok")

Complexity

Let n be the array length and k' = min(k, n) the number of workers actually used.

• partition: O(k') time, O(k') space for the bounds.
• Parallel phase: each worker sorts about n/k' elements in O((n/k') log(n/k')). Total work is O(n log(n/k')). With real parallelism the wall-clock time is O((n/k') log(n/k')), and thread start-up costs O(k').
• k_way_merge: O(n log k') time, O(k') heap space, plus O(n) for the output.
• Overall: O(n log n) work, O(n) extra space (the working copy and the output). On a single core the merge is the serial bottleneck, and it is O(n log k').

Test cases

• nums = [5, 1, 4, 2, 3], k = 2: segments [5,1,4], [2,3] → [1, 2, 3, 4, 5].
• nums = [9, 8, 7, 6, 5, 4, 3, 2, 1, 0], k = 3: segment sizes 4, 3, 3 → [0..9].
• nums = [], k = 3 → [] (no threads started).
• nums = [7], k = 4 → [7] (only one worker is used).
• k > n: nums = [2, 1], k = 5 → [1, 2], with segments (0,1), (1,2).
• k = 1 → a plain sort with no real merge.
• k = n → every segment holds one element, and the merge does all the work.
• Duplicates across segments: [3, 3, 1, 1, 2, 2], k = 3 → [1, 1, 2, 2, 3, 3].
• Negative numbers: [-2, 9, 0, -2, 5], k = 2 → [-2, -2, 0, 5, 9].
• Already sorted and reverse-sorted inputs with several values of k.
• Randomized check: compare against sorted(nums) for random arrays and random k.
• k <= 0 → ValueError.

Follow-ups

1. Parallelize the merge too. Merge the runs in a tree: in round j, worker pairs merge adjacent runs, giving log k rounds with a barrier between rounds. Each worker writes only to its target range of a second buffer, and the two buffers swap each round (ping-pong buffers). This removes the single-threaded O(n log k) tail. For full parallelism inside one merge, split it by co-ranking: binary-search the split points so that each thread produces a fixed output slice.
2. In place / no extra array. The constraint "each worker writes only its segment" holds during sorting. For the merge, discuss an in-place merge (rotations, O(n log n) per merge level, rarely worth it) against the standard trade-off of one O(n) auxiliary buffer.
3. Thread pool and reuse. Replace raw threads with a ThreadPoolExecutor or ProcessPoolExecutor, and wait on futures. In Python, use processes and shared memory (multiprocessing.shared_memory or array), because the GIL serializes CPU-bound threads. Also handle worker exceptions: future.result() re-raises them, so the caller shouldn't hang or get a partial result.
4. Data larger than memory (external sort). Each worker sorts a chunk and spills it to disk as a sorted run. The k-way heap merge then streams from buffered file readers. This is the same algorithm as above, now with I/O sizing (read buffers of size M/k).
5. Load balancing or skew. Evenly sized segments are not always equal work, for example with expensive comparators or nearly sorted chunks. Use more segments than threads with a work queue, or sample-sort: pick k-1 splitters, bucket each element so worker i owns a value range, and then the output is a concatenation with no merge at all.

Interviewer expectations

Mid-level (L4)

• Correct even partition, including k > n and the empty array.
• Starts the workers and joins them all before merging, and can explain why the join is needed.
• Merges with a heap instead of sorting again, and states O(n log k) for the merge.

Senior (L5)

• States the disjoint-segment invariant: no locks are needed because the write sets don't overlap, and the join supplies the memory visibility.
• Discusses the GIL in Python and the switch to processes, and handles worker exceptions.
• Implements or clearly designs the parallel tree merge with ping-pong buffers and per-round barriers.

Staff (L6+)

• Reasons about where the time goes: the serial merge tail limits speedup (Amdahl), thread start-up cost makes small n not worth parallelizing, and segment sizes should be cache-sized.
• Compares the alternatives: tree merge, merge-path or co-rank partitioning, sample sort, and external merge sort for data that doesn't fit in RAM.
• Covers production concerns: bounded pools, cancellation, deterministic and stable output, and benchmarking against a single-threaded sorted baseline.