S031 · Sorting — Merge, Quick, and When to Trust the Built-in

"""sorting.py — implementations + benchmarks for the three foundational sorts.
 
Run:  python sorting.py
"""
from __future__ import annotations
import heapq
import random
import time
 
 
# ------------------------------------------------------------------
# 1) Merge sort — the interview classic
# ------------------------------------------------------------------
def merge_sort(a: list[int]) -> list[int]:
    if len(a) <= 1:
        return a
    mid = len(a) // 2
    left = merge_sort(a[:mid])
    right = merge_sort(a[mid:])
    return _merge(left, right)
 
 
def _merge(left: list[int], right: list[int]) -> list[int]:
    out: list[int] = []
    i = j = 0
    while i < len(left) and j < len(right):
        if left[i] <= right[j]:
            out.append(left[i]); i += 1
        else:
            out.append(right[j]); j += 1
    out.extend(left[i:])
    out.extend(right[j:])
    return out
 
 
# ------------------------------------------------------------------
# 2) Quicksort — in-place with random pivot (avoids worst case on sorted input)
# ------------------------------------------------------------------
def quicksort(a: list[int], lo: int = 0, hi: int | None = None) -> None:
    if hi is None:
        hi = len(a) - 1
    if lo >= hi:
        return
    p = _partition(a, lo, hi)
    quicksort(a, lo, p - 1)
    quicksort(a, p + 1, hi)
 
 
def _partition(a: list[int], lo: int, hi: int) -> int:
    # Random pivot — SWAP to end. Prevents O(n^2) on already-sorted input.
    pivot_idx = random.randint(lo, hi)
    a[pivot_idx], a[hi] = a[hi], a[pivot_idx]
    pivot = a[hi]
    i = lo - 1
    for j in range(lo, hi):
        if a[j] <= pivot:
            i += 1
            a[i], a[j] = a[j], a[i]
    a[i + 1], a[hi] = a[hi], a[i + 1]
    return i + 1
 
 
# ------------------------------------------------------------------
# 3) Heap sort — build heap, repeatedly extract min
# ------------------------------------------------------------------
def heap_sort(a: list[int]) -> list[int]:
    heap = list(a)                   # copy so we don't mutate input
    heapq.heapify(heap)              # O(n)
    return [heapq.heappop(heap) for _ in range(len(heap))]   # n * O(log n) = O(n log n)
 
 
# ------------------------------------------------------------------
# 4) Insertion sort — perfect for small or almost-sorted lists
# ------------------------------------------------------------------
def insertion_sort(a: list[int]) -> list[int]:
    a = list(a)
    for i in range(1, len(a)):
        x = a[i]
        j = i - 1
        while j >= 0 and a[j] > x:
            a[j + 1] = a[j]
            j -= 1
        a[j + 1] = x
    return a
 
 
# ------------------------------------------------------------------
# 5) External merge — sorted chunks + k-way heap merge
# ------------------------------------------------------------------
def external_merge(chunks: list[list[int]]) -> list[int]:
    """Merge already-sorted chunks in O(n log k)."""
    return list(heapq.merge(*chunks))       # stdlib's k-way merge
 
 
# ------------------------------------------------------------------
# 6) Stability demo — TimSort preserves equal-key order; naive quicksort does not
# ------------------------------------------------------------------
def stability_demo() -> None:
    people = [("Alice", 30), ("Bob", 25), ("Carol", 30), ("Dan", 25), ("Eve", 30)]
    # Sort by age with sorted() — stable, preserves name order within each age group
    by_age_stable = sorted(people, key=lambda p: p[1])
    print("stable  (TimSort):", by_age_stable)
    # Now shuffle then sort — same output because we've re-sorted
    print("(if we sorted by NAME first then age, we'd get age-primary, name-secondary sorted)")
 
 
# ------------------------------------------------------------------
# 7) Benchmark our sorts vs Python's sorted() on 100k random ints
# ------------------------------------------------------------------
def benchmark() -> None:
    n = 100_000
    data = [random.randint(0, 1_000_000) for _ in range(n)]
 
    for name, fn in [
        ("merge_sort",  lambda d: merge_sort(d)),
        ("quicksort",   lambda d: (quicksort(d[:]), d)[0] or sorted(d)),  # returns None; use sorted for verify
        ("heap_sort",   lambda d: heap_sort(d)),
        ("sorted() TimSort", lambda d: sorted(d)),
    ]:
        d = data[:]
        t = time.perf_counter()
        if name == "quicksort":
            quicksort(d)
            result = d
        else:
            result = fn(d)
        dt = time.perf_counter() - t
        assert result == sorted(data)
        print(f"{name:20s} n={n} : {dt*1000:.1f} ms")
 
 
if __name__ == "__main__":
    print("merge_sort [3,1,4,1,5,9,2,6]:", merge_sort([3, 1, 4, 1, 5, 9, 2, 6]))
    a = [3, 1, 4, 1, 5, 9, 2, 6]; quicksort(a); print("quicksort in-place:", a)
    print("heap_sort  [3,1,4,1,5,9,2,6]:", heap_sort([3, 1, 4, 1, 5, 9, 2, 6]))
    print("insertion  [3,1,4,1,5,9,2,6]:", insertion_sort([3, 1, 4, 1, 5, 9, 2, 6]))
    print("external_merge:", external_merge([[1, 4, 7], [2, 5, 8], [3, 6, 9]]))
    stability_demo()
    benchmark()