Algorithm Search⌘K

Timsort in Python

Hybrid stable sort that merges detected runs with galloping and insertion sort; the default in Python and Java objects.

File
timsort.py, 294 lines
Level
I0 not checked yet
Toolchain
Python 3.14, pytest 9
Test framework
pytest
Run the tests
python3 -m pytest
timsort.pyRawGitHub
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/timsort/. The authors are listed in prov/PM6BZ3.jsonl.

from typing import Generic, Protocol, TypeVar


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


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

MIN_MERGE = 64
MIN_GALLOP = 7


def min_run(n: int) -> int:
    """Return the least run length: n itself if n < 64, else a value in [32, 64]."""
    r = 0
    while n >= MIN_MERGE:
        r |= n & 1
        n >>= 1
    return n + r


def count_run(a: list[T], lo: int, hi: int) -> int:
    """Return the length of the run that starts at lo. A strictly decreasing run is reversed."""
    if lo + 1 == hi:
        return 1
    i = lo + 1
    if a[i] < a[lo]:
        while i + 1 < hi and a[i + 1] < a[i]:
            i += 1
        a[lo : i + 1] = a[lo : i + 1][::-1]
    else:
        while i + 1 < hi and not a[i + 1] < a[i]:
            i += 1
    return i + 1 - lo


def binary_insertion_sort(a: list[T], lo: int, hi: int, start: int) -> None:
    """Sort a[lo:hi], given that a[lo:start] is sorted. Each item goes after its equals."""
    for i in range(start, hi):
        pivot = a[i]
        left, right = lo, i
        while left < right:
            mid = (left + right) // 2
            if pivot < a[mid]:
                right = mid
            else:
                left = mid + 1
        a[left + 1 : i + 1] = a[left:i]
        a[left] = pivot


def gallop_left(key: T, a: list[T], lo: int, n: int, hint: int) -> int:
    """Return k in [0, n] with a[lo + k - 1] < key <= a[lo + k], searching from lo + hint."""
    last, ofs = 0, 1
    if a[lo + hint] < key:
        maxofs = n - hint
        while ofs < maxofs and a[lo + hint + ofs] < key:
            last, ofs = ofs, 2 * ofs + 1
        ofs = min(ofs, maxofs)
        last, ofs = last + hint, ofs + hint
    else:
        maxofs = hint + 1
        while ofs < maxofs and not a[lo + hint - ofs] < key:
            last, ofs = ofs, 2 * ofs + 1
        ofs = min(ofs, maxofs)
        last, ofs = hint - ofs, hint - last
    last += 1  # now a[lo + last - 1] < key <= a[lo + ofs]
    while last < ofs:
        m = last + (ofs - last) // 2
        if a[lo + m] < key:
            last = m + 1
        else:
            ofs = m
    return ofs


def gallop_right(key: T, a: list[T], lo: int, n: int, hint: int) -> int:
    """Return k in [0, n] with a[lo + k - 1] <= key < a[lo + k], searching from lo + hint."""
    last, ofs = 0, 1
    if key < a[lo + hint]:
        maxofs = hint + 1
        while ofs < maxofs and key < a[lo + hint - ofs]:
            last, ofs = ofs, 2 * ofs + 1
        ofs = min(ofs, maxofs)
        last, ofs = hint - ofs, hint - last
    else:
        maxofs = n - hint
        while ofs < maxofs and not key < a[lo + hint + ofs]:
            last, ofs = ofs, 2 * ofs + 1
        ofs = min(ofs, maxofs)
        last, ofs = last + hint, ofs + hint
    last += 1  # now a[lo + last - 1] <= key < a[lo + ofs]
    while last < ofs:
        m = last + (ofs - last) // 2
        if key < a[lo + m]:
            ofs = m
        else:
            last = m + 1
    return ofs


class TimSort(Generic[T]):
    def __init__(self, a: list[T]) -> None:
        self.a = a
        self.min_gallop = MIN_GALLOP
        self.runs: list[tuple[int, int]] = []  # (base, length), from the bottom of the stack
        self.merges: list[tuple[int, int]] = []  # the lengths of each merge, for tests

    def merge_at(self, i: int) -> None:
        """Merge the runs i and i + 1 of the stack."""
        a = self.a
        base_a, len_a = self.runs[i]
        base_b, len_b = self.runs[i + 1]
        self.merges.append((len_a, len_b))
        self.runs[i] = (base_a, len_a + len_b)
        del self.runs[i + 1]
        # Items of A that are at most B[0] are already in place.
        k = gallop_right(a[base_b], a, base_a, len_a, 0)
        base_a += k
        len_a -= k
        if len_a == 0:
            return
        # Items of B that are at least the last item of A are already in place.
        len_b = gallop_left(a[base_a + len_a - 1], a, base_b, len_b, len_b - 1)
        if len_b == 0:
            return
        if len_a <= len_b:
            self.merge_lo(base_a, len_a, base_b, len_b)
        else:
            self.merge_hi(base_a, len_a, base_b, len_b)

    def merge_lo(self, base_a: int, len_a: int, base_b: int, len_b: int) -> None:
        """Merge from the left, with A copied to a buffer. On a tie the item of A goes first."""
        a = self.a
        tmp = a[base_a : base_a + len_a]
        i, j, dest = 0, base_b, base_a
        end_b = base_b + len_b
        while i < len_a and j < end_b:
            count_a = count_b = 0
            while i < len_a and j < end_b:  # one item at a time
                if a[j] < tmp[i]:
                    a[dest] = a[j]
                    j += 1
                    count_b += 1
                    count_a = 0
                    if count_b >= self.min_gallop:
                        break
                else:
                    a[dest] = tmp[i]
                    i += 1
                    count_a += 1
                    count_b = 0
                    if count_a >= self.min_gallop:
                        break
                dest += 1
            else:
                break
            dest += 1
            self.min_gallop += 1
            while i < len_a and j < end_b:  # galloping mode
                self.min_gallop -= self.min_gallop > 1
                k = gallop_right(a[j], tmp, i, len_a - i, 0)
                a[dest : dest + k] = tmp[i : i + k]
                dest += k
                i += k
                if i == len_a:
                    break
                a[dest] = a[j]
                dest += 1
                j += 1
                if j == end_b:
                    break
                k2 = gallop_left(tmp[i], a, j, end_b - j, 0)
                a[dest : dest + k2] = a[j : j + k2]
                dest += k2
                j += k2
                if j == end_b:
                    break
                a[dest] = tmp[i]
                dest += 1
                i += 1
                if i == len_a:
                    break
                if k < MIN_GALLOP and k2 < MIN_GALLOP:
                    self.min_gallop += 1
                    break
        a[dest : dest + len_a - i] = tmp[i:]

    def merge_hi(self, base_a: int, len_a: int, base_b: int, len_b: int) -> None:
        """Merge from the right, with B copied to a buffer. On a tie the item of B goes last."""
        a = self.a
        tmp = a[base_b : base_b + len_b]
        i, j = base_a + len_a - 1, len_b - 1  # the last items of A and of the buffer
        dest = base_b + len_b - 1
        while i >= base_a and j >= 0:
            count_a = count_b = 0
            while i >= base_a and j >= 0:  # one item at a time
                if tmp[j] < a[i]:
                    a[dest] = a[i]
                    i -= 1
                    count_a += 1
                    count_b = 0
                    if count_a >= self.min_gallop:
                        break
                else:
                    a[dest] = tmp[j]
                    j -= 1
                    count_b += 1
                    count_a = 0
                    if count_b >= self.min_gallop:
                        break
                dest -= 1
            else:
                break
            dest -= 1
            self.min_gallop += 1
            while i >= base_a and j >= 0:  # galloping mode
                self.min_gallop -= self.min_gallop > 1
                rem_a = i - base_a + 1
                k = rem_a - gallop_right(tmp[j], a, base_a, rem_a, rem_a - 1)
                a[dest - k + 1 : dest + 1] = a[i - k + 1 : i + 1]
                dest -= k
                i -= k
                if i < base_a:
                    break
                a[dest] = tmp[j]
                dest -= 1
                j -= 1
                if j < 0:
                    break
                k2 = j + 1 - gallop_left(a[i], tmp, 0, j + 1, j)
                a[dest - k2 + 1 : dest + 1] = tmp[j - k2 + 1 : j + 1]
                dest -= k2
                j -= k2
                if j < 0:
                    break
                a[dest] = a[i]
                dest -= 1
                i -= 1
                if i < base_a:
                    break
                if k < MIN_GALLOP and k2 < MIN_GALLOP:
                    self.min_gallop += 1
                    break
        a[base_a : base_a + j + 1] = tmp[: j + 1]

    def merge_collapse(self) -> None:
        """Restore the invariants on the run lengths, with the check of de Gouw et al. (2015)."""
        r = self.runs
        while len(r) > 1:
            n = len(r) - 2
            if (n > 0 and r[n - 1][1] <= r[n][1] + r[n + 1][1]) or (
                n > 1 and r[n - 2][1] <= r[n - 1][1] + r[n][1]
            ):
                if r[n - 1][1] < r[n + 1][1]:
                    n -= 1
                self.merge_at(n)
            elif r[n][1] <= r[n + 1][1]:
                self.merge_at(n)
            else:
                break

    def merge_force_collapse(self) -> None:
        r = self.runs
        while len(r) > 1:
            n = len(r) - 2
            if n > 0 and r[n - 1][1] < r[n + 1][1]:
                n -= 1
            self.merge_at(n)


def timsort(a: list[T]) -> list[tuple[int, int]]:
    """Sort the list a in place, stably. Return the lengths of the runs that were merged, in order."""
    n = len(a)
    ts = TimSort(a)
    if n < 2:
        return ts.merges
    minrun = min_run(n)  # @line L1
    lo = 0  # @line L2
    while lo < n:  # @line L3
        run = count_run(a, lo, n)  # @line L4
        if run < minrun:  # @line L5
            force = min(minrun, n - lo)
            binary_insertion_sort(a, lo, lo + force, lo + run)
            run = force
        ts.runs.append((lo, run))  # @line L6
        lo += run
        ts.merge_collapse()  # @line L7
    ts.merge_force_collapse()  # @line L8
    return ts.merges
test_timsort.py 94 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/timsort/. The authors are listed in prov/PM6BZ3.jsonl.

import bisect
import random

from timsort import binary_insertion_sort, count_run, gallop_left, gallop_right, min_run, timsort


class Key:
    __slots__ = ("k", "tag")

    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


def test_min_run() -> None:
    assert [min_run(n) for n in (1, 63, 64, 65, 100, 350, 2000, 10**5, 10**6)] == [1, 63, 32, 33, 50, 44, 63, 49, 62]
    for n in range(64, 5000):
        m = min_run(n)
        assert 32 <= m <= 64


def test_example_merges() -> None:
    a: list[int] = []
    base = 10**6
    for length in (120, 80, 60, 50, 40):
        base -= 1000
        a += list(range(base, base + length))
    expected = sorted(a)
    assert timsort(a) == [(80, 60), (120, 140), (50, 40), (260, 90)]
    assert a == expected


def test_count_run() -> None:
    a = [5, 4, 3, 3, 1]
    assert count_run(a, 0, 5) == 3
    assert a[:3] == [3, 4, 5]
    b = [1, 2, 2, 0]
    assert count_run(b, 0, 4) == 3


def test_binary_insertion_sort_is_stable() -> None:
    a = [Key(2, 0), Key(1, 1), Key(2, 2), Key(1, 3)]
    binary_insertion_sort(a, 0, 4, 1)
    assert [(x.k, x.tag) for x in a] == [(1, 1), (1, 3), (2, 0), (2, 2)]


def test_gallops_match_bisect() -> None:
    rng = random.Random(1)
    for _ in range(2000):
        a = sorted(rng.randint(0, 9) for _ in range(rng.randint(1, 20)))
        pre = [rng.randint(0, 5) for _ in range(rng.randint(0, 3))]
        b = pre + a
        key = rng.randint(-1, 10)
        hint = rng.randrange(len(a))
        assert gallop_left(key, b, len(pre), len(a), hint) == bisect.bisect_left(a, key)
        assert gallop_right(key, b, len(pre), len(a), hint) == bisect.bisect_right(a, key)


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


def patterned(rng: random.Random, n: int, mode: int) -> list[int]:
    if mode == 0:
        return [rng.randint(0, 3) for _ in range(n)]
    if mode == 1:
        return [rng.randint(0, 10**6) for _ in range(n)]
    if mode == 2:
        ks: list[int] = []
        while len(ks) < n:
            length = rng.randint(1, 200)
            block = [rng.randint(0, 100) for _ in range(length)]
            ks += sorted(block, reverse=rng.random() < 0.3)
        return ks[:n]
    half = n // 2
    return sorted(rng.randint(0, 100) for _ in range(half)) + sorted(rng.randint(0, 100) for _ in range(n - half))


def test_stable_on_patterns() -> None:
    rng = random.Random(2)
    for _ in range(300):
        n = rng.choice([5, 63, 64, 65, 300, 1000])
        ks = patterned(rng, n, rng.randrange(4))
        b = [Key(k, i) for i, k in enumerate(ks)]
        timsort(b)
        assert [(x.k, x.tag) for x in b] == sorted((k, i) for i, k in enumerate(ks))

What the next level needs

  • the tests did not run in this build

Read the full page about Timsort