Tree sort in Python

Insert all keys into a binary search tree, then read them in order; Theta(n log n) average.

File: tree_sort.py, 70 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/tree-sort/. The authors are listed in prov/P5E3N3.jsonl.

from typing import Generic, Protocol, TypeVar


class Ordered(Protocol):
    def __lt__(self, other: "Ordered", /) -> bool: ...


T = TypeVar("T", bound=Ordered)


class Node(Generic[T]):
    __slots__ = ("key", "left", "right")

    def __init__(self, key: T) -> None:
        self.key = key
        self.left: "Node[T] | None" = None
        self.right: "Node[T] | None" = None


def insert(root: "Node[T] | None", x: T) -> tuple["Node[T]", int]:
    """Insert x without recursion: left if x < key, else right. Return the root and the number of comparisons."""
    node = Node(x)
    if root is None:
        return node, 0
    u = root
    count = 0
    while True:
        count += 1
        if x < u.key:
            if u.left is None:
                u.left = node
                return root, count
            u = u.left
        else:
            if u.right is None:
                u.right = node
                return root, count
            u = u.right


def inorder(root: "Node[T] | None") -> list[T]:
    """Return the keys in inorder, with an explicit stack."""
    out: list[T] = []
    stack: list[Node[T]] = []
    u = root
    while stack or u is not None:
        while u is not None:
            stack.append(u)
            u = u.left
        u = stack.pop()
        out.append(u.key)
        u = u.right
    return out


def tree_sort(a: list[T]) -> int:
    """Sort the list a in place, stably. Return the number of comparisons."""
    root: Node[T] | None = None  # @line L1
    total = 0
    for x in a:  # @line L2
        root, c = insert(root, x)  # @line L3
        total += c
    i = 0  # @line L4
    for key in inorder(root):  # @line L5
        a[i] = key  # @line L6
        i += 1
    return total

test_tree_sort.py 59 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/tree-sort/. The authors are listed in prov/P5E3N3.jsonl.

import random

from tree_sort import inorder, insert, tree_sort


def test_example() -> None:
    a = [5, 3, 8, 1, 4, 7, 9, 2, 6]
    assert tree_sort(a) == 16
    assert a == list(range(1, 10))


def test_sorted_input_is_quadratic() -> None:
    a = list(range(10))
    assert tree_sort(a) == 45
    a = list(range(3000))
    assert tree_sort(a) == 3000 * 2999 // 2


def test_inorder() -> None:
    root = None
    for x in [2, 1, 3]:
        root, _ = insert(root, x)
    assert inorder(root) == [1, 2, 3]
    assert inorder(None) == []


def test_edges() -> None:
    for a in ([], [1], [2, 1], [1, 1, 1], [3, 2, 1]):
        b = list(a)
        tree_sort(b)
        assert b == sorted(a)


def test_stable() -> None:
    class Key:
        def __init__(self, k: int, tag: int) -> None:
            self.k, self.tag = k, tag

        def __lt__(self, other: "Key") -> bool:
            return self.k < other.k

    rng = random.Random(4)
    for _ in range(200):
        pairs = [(rng.randint(0, 4), i) for i in range(rng.randint(0, 30))]
        b = [Key(k, t) for k, t in pairs]
        tree_sort(b)
        assert [(x.k, x.tag) for x in b] == sorted(pairs)


def test_random_against_sorted() -> None:
    rng = random.Random(9)
    for _ in range(100):
        a = [rng.randint(-50, 50) for _ in range(rng.randint(0, 200))]
        b = list(a)
        tree_sort(b)
        assert b == sorted(a)

What the next level needs

the tests did not run in this build

Read the full page about Tree sort