American flag sort in Python
In-place MSD radix sort that permutes keys into digit bins by cycle-following, named for the Dutch flag problem.
- File
american_flag_sort.py, 48 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/american-flag-sort/. The authors are listed in prov/GW15GA.jsonl.
def _c(s: str, d: int) -> int:
"""The bucket of s at depth d: 0 when s has ended, else the code of s[d] plus 1."""
return ord(s[d]) + 1 if d < len(s) else 0
def american_flag_sort(a: list[str], r: int = 256) -> list[str]:
"""Sort the list a of strings in lexicographic order, in place, and return it.
Precondition: each character has a code less than r. The sort is not stable.
The cost is O(r (n + D)) time, where D is the total length of the
distinguishing prefixes, and O(r) counters per step with no item buffer.
"""
for s in a:
for ch in s:
if ord(ch) >= r:
raise ValueError("a character is outside the alphabet")
stack = [(0, len(a), 0)] # @line L1
while stack: # @line L2
lo, hi, d = stack.pop() # @line L3
if hi - lo <= 1: # @line L4
continue
count = [0] * (r + 1) # @line L5
for i in range(lo, hi):
count[_c(a[i], d)] += 1
start = [0] * (r + 1) # @line L6
pos = lo
for b in range(r + 1):
start[b] = pos
pos += count[b]
nxt = start[:]
end = [start[b] + count[b] for b in range(r + 1)]
for cur in range(r + 1): # @line L7
while nxt[cur] < end[cur]: # @line L8
b = _c(a[nxt[cur]], d) # @line L9
if b == cur: # @line L10
nxt[cur] += 1
else: # @line L11
j = nxt[b]
a[nxt[cur]], a[j] = a[j], a[nxt[cur]]
nxt[b] = j + 1
for b in range(r, 0, -1): # @line L12
if count[b] >= 2:
stack.append((start[b], end[b], d + 1))
return a # @line L13
test_american_flag_sort.py 47 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/american-flag-sort/. The authors are listed in prov/GW15GA.jsonl.
import random
import unittest
from american_flag_sort import american_flag_sort
class TestAmericanFlagSort(unittest.TestCase):
def test_page_example(self):
a = "she sells sea shells by the sea shore".split()
self.assertEqual(american_flag_sort(a), ["by", "sea", "sea", "sells", "she", "shells", "shore", "the"])
def test_in_place(self):
a = ["c", "a", "b"]
self.assertIs(american_flag_sort(a), a)
self.assertEqual(a, ["a", "b", "c"])
def test_edges(self):
self.assertEqual(american_flag_sort([]), [])
self.assertEqual(american_flag_sort(["only"]), ["only"])
self.assertEqual(american_flag_sort(["", "b", "", "a"]), ["", "", "a", "b"])
self.assertEqual(american_flag_sort(["ab", "b", "abc", "a", "ba"]), ["a", "ab", "abc", "b", "ba"])
self.assertEqual(american_flag_sort(["same"] * 9), ["same"] * 9)
with self.assertRaises(ValueError):
american_flag_sort(["Ā"])
def test_small_alphabet(self):
a = ["10", "0", "1", "01", "", "00"]
self.assertEqual(american_flag_sort(list(a), r=50), sorted(a))
def test_long_common_prefix(self):
p = "k" * 4000
a = [p + "z", p + "a", p, p + "m"]
self.assertEqual(american_flag_sort(a), sorted(a))
def test_random(self):
rng = random.Random(1993)
for _ in range(600):
n = rng.randint(0, 50)
k = rng.choice([1, 2, 5, 26])
a = ["".join(chr(97 + rng.randrange(k)) for _ in range(rng.randint(0, 6))) for _ in range(n)]
self.assertEqual(american_flag_sort(list(a)), sorted(a))
if __name__ == "__main__":
unittest.main()
What the next level needs
- the tests did not run in this build