Data Structures

To my surprise, a fixed-size array can live on the stack and be passed around into sub-functions. You must only remember to never use it after the defined function. When you pass it, callers just get a pointer to where it lives on your stack frame.

I added some basic push() and pop() functions to see what mutating it and passing around metadata feels like. This gives it lite stack (data structure) semantics.

#include <stdio.h>

/**
 * Returns new len.
 * prereq: 0 <= len <= capacity
 */
int push(int *a, int capacity, int len, int value) {
    if (len == capacity) {
        printf("Warning: Can't push %d: array is at capacity (%d).\n", value, capacity);
        return len;
    }
    a[len] = value;
    return len + 1;
}

/**
 * Returns new len.
 * prereq: 0 <= len <= capacity
 */
int pop(int *a, int len) {
    if (len == 0) {
        printf("Warning, can't pop: array is empty.\n");
        return len;
    }
    // just leave the value there! :-)
    return len - 1;
}

/**
 * Prints `a`'s values in [0, len).
 * prereq: 0 <= len <= capacity
*/
void display(int *a, int capacity, int len) {
    printf("int* array at %p, capacity = %d, len = %d\n", (void*)a, capacity, len);
    for (int i = 0; i < len; i++) {
        printf("a[%d] = %d\n", i, a[i]);
    }
    printf("\n");
}

/**
 * Fill all of `a` with 0.
 */
void zero(int *a, int capacity) {
    for (int i = 0; i < capacity; i++) {
        a[i] = 0;
    }
}

int main() {
    // dynamic stack array
    int capacity = 10;
    int len = 0;
    int a[capacity]; // can't init (= {0}) here as dynamic
    // zero(a, capacity); // instead. (optional.)

    display(a, capacity, len); // nothing

    for (int i = 0; i < capacity + 1; i++) {
        len = push(a, capacity, len, 100 + i);
    }

    display(a, capacity, len); // 100 -- 109

    for (int j = 0; j < 3; j++) {
        len = pop(a, len);
    }

    display(a, capacity, len); // 100 - 106

    return 0;
}

Doing it as a raw array on the stack is fun to immediately see pain points:

• Not being able to expand the array. (Allocating heap memory so it’s expandable makes sense.)
• Having to pass all metadata as function args is cumbersome. (Structs make sense.)

First, we’ll assume the existence of a good hash function and compression function:

• a hash function h maps any item I we’d like to store into, let’s say, an int H within a fixed range of nonnegative integers [0, M] in Z. (E.g., 32-bit or 64-bit integers.)
  • a good hash function makes it unlikely that two distinct items I and I’ produce the same hash H while there are < M items; i.e., hashes for input objects are distributed uniformly across [0, M]
• a compression function (AKA “index mapping”) maps [0, M] to the size of our actual backing key array of length K.
  • Hmm, trivially this could be done by simply mod, with our array key H’ = H % K
  • However, there exist good hash functions that become in practice bad because they collide after this mod
  • Studied strategies (from looking up; would need to read more)
    • division: H mod K, but choose K to so it’s unlikely to expose patterns in hash values (e.g., prime K away from power of 2). However, this seems bad because your backing away now won’t be near a power of two, which is good for memory alignment?
    • multiplication: choose some constant e in (0, 1), do H * e, take the fraction bits f, then do f * K
    • MAD (multiply-add-divide): (aH + b) mod p mod K (with p a prime > K)

I’m thinking this is backed by an array because an array gives us O(1) access once we’ve gotten the compressed hash.

I understand why hashes need to be stable. If an item’s hash changes, you wouldn’t be able to locate it in the map anymore.

Now, we can take any item we’d like to store and produce an index in [0, K) of a backing array. Each entry in the array could store:

• the item itself (open addressing; need probing)
  • a primitive if fixed length (e.g., <= 64 bits)
  • the full item (will need space + copy)
  • primitive or pointer (e.g., still always <= 64 bits, but potential optimization for primitive items)
  • always pointers (simplest, natural for many langs)
• a (pointer to a) data structure (separate chaining; for collisions)
  • dynamic array
  • linked list
  • another hash map (len K’, hopefully to avoid collisions with K)
    • The sub-hash-map strategy seems sus. Feels like you should just grow the original hash map once there’s enough pressure. Though I guess could help if you have a pathologically bad initial H->K, then a great H->K’?
  • between dynamic array and linked list, I like linked list. You’ll have O(n) retrieval and O(1) insertions at end regardless, so might as well have O(1) deletion rather O(n)

So far I’ve been assuming we want to store and retrieve an item I. This seems suitable to implement a set: O(1) lookup and retrieval. I think we also need an equality check for items, because after a collision we need to know whether the candidate and existing entry match.

However, we may want to hash user-provided key-value pairs (K, V), to implement a mapping. In this case, I think we simply hash only the key, and explicitly store the key alongside the value in the data structure so that we can explicitly use the key

Growing. If we track the number of elements in the data structure, we can decide to grow the data structure at some fullness threshold. (“Pressure?”) Naively, we’d

• choose a new K’ > K
• iterate over all items, and find their new index in K’
  • if we also stored each item’s hash (which should be OK b/c stable), we wouldn’t have to re-hash. This feels like a minor optimization but could be good if hashing is expensive, like checksumming a lot of data
  • regardless, I think we’d need to re-run the compression function
    • we could skip this if we have an equation to transform a compressed value from an input of K to K’
    • however, this is probably as fast as running the compression function again, because they seem simple
  • then, we’d need to re-insert freshly into the new map
    • if we have strong guarantees about the ordering of items under some K -> K’ assumptions (i.e., haven’t undergone a mod), AND we don’t have any collisions (no open addressing or chaining happened), we could potentially directly copy chunks of items
    • … but this seems like an over-optimization that would be useful rarely. zero collisions before growing seems unlikely. and it seems like the chained data structures might need some new bookkeeping.
    • in fact, the whole reason we’re growing the data structure in the first place implies we likely have collisions, which means naively copying is probably a bad idea
• this is another O(n) operation that I think amortizes well, because we tie resizes to inserts and use a multiplicative growth strategy.

Iteration.

• If we have only a backing array of len K, and we have n stored items with n << K, then O(K) dominates the iteration time. (We iterate over all buckets.) Once n >= K, then O(n) dominates. (I’m not sure how to write this in big-O.)
• We lack the ability to provide, e.g., iteration order by insertion
• To provide insertion order and strict O(n) iteration, we could keep another data structure. (Gets (value access) always O(1) from hash map.)
  • We could use a dynamic array that we push new insertions to at the end.
    • O(n) deletion
      • shift and copy, same as array
      • this is bad, because it ruins our O(1) hash map deletion
    • O(1) get by index (what’s the 4th item?)
    • interestingly: O(n) by value’s index (when was this value inserted?)
      • Unless we kept the value’s index in the hash table as well, but this would require O(n) updates on every deletion (which is already true…), and if we did this might as well skip the array entirely
  • We could use a linked list
    • O(n) deletion because it’s O(n) retrieval (where is node on linked list?)
      • Fast list position retrieval could again be circumvented by again the position in the node, but would again require O(n) updates on deletion because of bookkeeping
    • O(n) get by index (could be non-goal)
  • What if we used a doubly linked list within the hash map entries?
    • Keep pointers to first and last items
    • O(1) deletion: find and patch up neighbors’ links
    • still O(1) insertion: added to end
    • O(n) by index (4th item?) or value index (which item?)
• Python’s uses tombstones in the backing dynamic array (which also holds the data!).
  • This is interesting because it, and the backing doubly-linked list, try to match the time complexity of a hash map O(1) get / insert / delete, and support easy forward / backward iteration, but lack O(1) lookup by insertion index

Handling open addressing Open addressing I believe means upon collision you walk forward until you find an opening. I intuitively have never liked this because it makes the top level data structure messier.

• Naively, retrievals become “hit then walk”
• But the issue is deletions can leave holes
  • To overcome the holes, we need some bookkeeping strategy. We could track the maximum we’ve ever walked, or a per-bucket max walk. Then always walk that far. Per-bucket seems better because you can actually decrement it later upon deletions (though maybe over-optimization to do so).
• Actually, it seems like leaving holes isn’t commonly done. The following strategies assume no holes.
• Tombstones (lazy deletion): keep entries as “occupied” so searching forwards keeps looking until a hole. This is cool because insertions can still use tombstones.
• Shifting (more typical?): when you’d delete, shift things back to avoid any holes.
  • This seems… really complicated? Couldn’t you have to shift n entries, hopping over entries and clusters from other hashes? I think yes. But I guess you do have the nice invariant of “no holes within a cluster,” so as soon as you reach a hole you’re done.

Probing strategies

• linear probing (walk 1, clustering, great cache perf)
• double hashing (walk deterministic n computed from bucket (e.g., index), no clustering, bad cache perf)
• quadratic probing (walk 1,2,3,…; moderate clustering & caching)

Because a binary heap implementation isn’t completely trivial, it’s worth sketching out.

The core logic of the operations is quite compact. Unfortunately, basic comments and invalid arg checking easily doubles the size.

"""
Binary (max, WLOG) heap, implements a priority queue.

throughout:
    v = value
    i = index

Storing `int`s as values directly for simplicity.

If this were tracking an order on actual objects, I think there'd be three main
approaches:

    1. store objects themselves in the backing array. access `.priority` to compare,
       swap whole objects

    2. store `{priority, pointer}` in the backing array. access `.priority` to compare,
       swap just mini structs.

    3. store indices into an array `a` of objects in the backing array. `a[i].priority`
       to compare, swap just ints.

(I think 3. is probably the most elegant.)
"""

import random

# index operations for backing the binary heap with an array.
# existence of indices not guaranteed!


def parent(i: int) -> int:
    return (i - 1) // 2  # floor


def children(i: int) -> tuple[int, int]:
    return (2 * i + 1, 2 * i + 2)


# normally not a big class guy, but it beats passing the data array everywhere


class BinaryMaxHeap:
    def __init__(self):
        # dynamic array, handles resizing. thanks python
        self.values: list[int] = []

    def _bubble_up(self, i: int):
        """Recursively swap i with its parent if i is bigger.

        AKA up-heap, percolate-up, sift-up, trickle-up, swim-up, heapify-up,
        cascade-up, fix-up

        O(log n)
        """
        if i == 0:
            return  # at root

        p = parent(i)
        if self.values[i] > self.values[p]:
            self.values[i], self.values[p] = self.values[p], self.values[i]
            self._bubble_up(p)

    def _bubble_down(self, i: int):
        """Recursively swap i with its largest child if either is bigger than it.

        AKA down-heap, percolate-down, sift-down, sink-down, trickle down,
        heapify-down, cascade-down, fix-down, extract-min or extract-max, or
        just heapify

        O(log n)
        """
        c1, c2 = children(i)
        candidates = [c for c in [c1, c2] if c < len(self.values)]  # validate
        if len(candidates) == 0:
            return  # at leaf
        max_val = max(self.values[c] for c in candidates)
        if max_val > self.values[i]:
            max_c = [c for c in candidates if self.values[c] == max_val][0]
            self.values[i], self.values[max_c] = (
                self.values[max_c],
                self.values[i],
            )
            self._bubble_down(max_c)

    # operations
    #
    # going with ValueError because I usually have invalid operations return
    # None, so trying exceptions instead as something new.

    def find_max(self):
        """O(1)"""
        if len(self.values) == 0:
            raise ValueError()
        return self.values[0]  # it's me, Max. I'm always right here.

    def insert(self, v: int):
        """O(log n)"""
        self.values.append(v)
        self._bubble_up(len(self.values) - 1)  # the inserted index

    def size(self) -> int:
        """Only seemed polite since invalid requests raise ValueError.

        O(1)
        """
        return len(self.values)

    def extract(self) -> int:
        """AKA "delete max;" remove and return the max val, O(log n)"""
        if len(self.values) == 0:
            raise ValueError()
        if len(self.values) == 1:
            return self.values.pop()

        ret = self.values[0]  # should I call find_max()?
        self.values[0] = self.values.pop()  # end --> root
        self._bubble_down(0)
        return ret

    def insert_then_extract(self, v: int) -> int:
        """insert a new value and extract the new largest. O(log n)"""
        # if what we're inserting is the new max, we can just return it
        if len(self.values) == 0 or v > self.values[0]:
            return v

        # else, we return the root, and bubble down
        ret = self.values[0]
        self.values[0] = v
        self._bubble_down(0)
        return ret

    def delete(self, i: int) -> int:
        """deletes value at index i and returns it. O(log n)"""
        if i >= len(self.values):
            raise ValueError()
        deleted = self.values[i]
        self.values[i] = self.values.pop()  # last element -> here
        if self.values[i] > deleted:
            # the last element is bigger than what was here, so we may
            # need to move it up
            self._bubble_up(i)
        else:
            # smaller, so may need to move it down
            self._bubble_down(i)
        return deleted

    def search(self, v: int) -> int:
        """Finds the first index i that has the value v, or ValueError if doesn't exist.

        O(n)

        Felt right to add this since so many functions are passed an index i,
        but practically it seems like you'd want to find the index i based on
        a value v."""
        return self.values.index(v)

    def decrease_key(self, i: int, v: int):
        """Decreases the value at index i to be v."""
        if i >= len(self.values):
            raise ValueError()
        self.values[i] = v
        self._bubble_down(i)  # as smaller, maybe move down

    def increase_key(self, i: int, v: int):
        """Increases the value at index i to be v."""
        if i >= len(self.values):
            raise ValueError()
        self.values[i] = v
        self._bubble_up(i)  # as larger, maybe move up

    def modify_key(self, i: int, v: int):
        """Modifies the value at index i to be v.

        Seems like you'd just want this as an API right?
        """
        if i >= len(self.values):
            raise ValueError()
        if v < self.values[i]:
            self.decrease_key(i, v)
        else:
            self.increase_key(i, v)

    def meld(self, other: "BinaryMaxHeap"):
        """AKA 'merge', combines `other` (m) into this (n) heap.

        Does not mutate `other`.

        Binary heaps have no special meld runtime. The two main options are:
            1. insert each of m          - O(m log(n + m))
            2. heapify both from scratch - O(n + m)

        Here we simply do 1.
        """
        # copy values first in case we're doing `h.meld(h)`
        for v in other.values[:]:
            self.insert(v)

    @staticmethod
    def build_heap(vs: list[int]) -> "BinaryMaxHeap":
        """Constructs a heap from provided values vs.

        Rather than inserting n elements (O(n log n), Williams), we
        immediately construct a binary tree (shape property) then order it
        (heap property) (O(n), Floyd).

        Running time is O(n) due to a complicated analysis. Naively we're
        skipping the leaves (~n/2) and doing h (log n) operations per node,
        but this ends up being good enough to reduce us from O(n log n) to
        O(n).
        """
        h = BinaryMaxHeap()
        h.values = vs[:]  # copy over. complete binary tree achieved!

        if len(h.values) == 0:
            return h  # avoid no values edge case range would hit

        # now, make sure every node is bigger than its children, or
        # moves down if not. trivially true for leaf nodes, so start at
        # penultimate level and work back up to root node, inclusive.
        for i in range(len(h.values) // 2, -1, -1):
            h._bubble_down(i)

        return h


def extract_all(h: BinaryMaxHeap) -> list[int]:
    """Calls extract() on h until it's empty. Returns list of results."""
    ordered = []
    while h.size() > 0:
        ordered.append(h.extract())
    return ordered


def main():
    # build from shuffled [1, 10]
    vs = list(range(1, 11))
    random.shuffle(vs)
    h = BinaryMaxHeap.build_heap(vs)
    ordered = extract_all(h)
    print("original:    ", vs)
    print("max heap'd:  ", ordered)

    # build again, this time from inserts, and merge
    for v in vs:
        h.insert(v)
    h.meld(h)
    print("self-melded: ", extract_all(h))

    # try more operations
    h = BinaryMaxHeap.build_heap(vs)
    h.delete(0)  # delete 10
    h.decrease_key(0, 2)  # 9 -> 2
    h.modify_key(0, 2)  # 8 -> 2
    h.increase_key(h.search(1), 2)  # 1 -> 2
    print("7-3, 2 2 2 2:", extract_all(h))


if __name__ == "__main__":
    main()

From my source repo.

"""Super basic BST, no balancing, just to get the BST flavor.

Had basically all functional, migrated more to being instance methods, but a few
functional-ish ones remain (_traverse(), _search(), maybe _replace() counts too).
"""

import random
from typing import Optional


class BinaryTreeNode:
    def __init__(
        self,
        value: int,
        parent: Optional["BinaryTreeNode"],
        left: Optional["BinaryTreeNode"],
        right: Optional["BinaryTreeNode"],
    ):
        self.value = value
        self.parent = parent
        self.left = left
        self.right = right

    def __repr__(self) -> str:
        return f"({self.value})"

    def insert(self, new: "BinaryTreeNode"):
        """inserts new into self.

        i preferred the non-class more recursive version from before, with the pattern:
            n.left = insert(n.left, new)
            n.left.parent = n
            return n

        ... but this convention only makes sense if you recurse down to calling on None
        and just returning new, which you can't do as an instance method. so making
        this an instance method meant changing this to the more branching way.
        """
        # equal; shouldn't happen
        if self.value == new.value:
            print("Error: assuming unique values. Not inserting.")

        if new.value < self.value:
            if self.left is not None:
                self.left.insert(new)
            else:
                self.left = new
                new.parent = self
        else:
            if self.right is not None:
                self.right.insert(new)
            else:
                self.right = new
                new.parent = self

    def minimum(self: "BinaryTreeNode") -> "BinaryTreeNode":
        """Traverse to leftmost leaf to find minimum."""
        # I kept the `self is None`` case in the BST's call, because we don't
        # actually want to recurse until None, we want to stop at a leaf.
        if self.left is None:
            return self
        return self.left.minimum()

    def maximum(self: "BinaryTreeNode") -> "BinaryTreeNode":
        """Traverse to rightmost leaf to find maximum node."""
        if self.right is None:
            return self
        return self.right.maximum()

    def successor(self: "BinaryTreeNode") -> Optional["BinaryTreeNode"]:
        """returns the node with the smallest value > the self's value, None if max."""
        # the logic here was quite tricky to convince myself of. x's successor is either:
        # 1. _minimum(x.right) if x.right exists, else
        # 2. move upwards until you go up-right (↗) once, then return that
        if self.right is not None:
            return self.right.minimum()

        cur, parent = self, self.parent
        while parent is not None and parent.right is cur:
            cur = parent
            parent = parent.parent
        return parent

    def predecessor(self: "BinaryTreeNode") -> Optional["BinaryTreeNode"]:
        """returns the node with the greatest value < the self's value, None if min."""
        # x's predecessor is either:
        # 1. _maximum(x.left) if x.left exists, else
        # 2. move upwards until you go up-left (↖) once, then return that
        if self.left is not None:
            return self.left.maximum()

        cur, parent = self, self.parent
        while parent is not None and parent.left is cur:
            cur = parent
            parent = parent.parent
        return parent

    def delete(self: "BinaryTreeNode", node: "BinaryTreeNode"):
        """
        Returns this subtree

        Assumes:
        - NO MORE node is not self (already checked in BST not root)
           - --> NO MORE node has a parent
        - node is in self's tree
        """
        n_children = len([c is None for c in [node.left, node.right]])

        # three cases:
        # (1) if node is a leaf, it's simply replaced by None
        if n_children == 0:
            _replace(node, None)
            return self if node is not self else None

        # (2) node has only one child. it's simply replaced by its child.
        # (makes sense because any values that would have been correctly
        # routed to it are fine being routed to its only child, because its
        # only child obeys the order node had with its parent.)
        if n_children == 1:
            child = node.left if node.left is not None else node.right
            _replace(node, child)
            return self if node is not self else child

        # (3) node has two children. we'll use its successor S to replace it.
        # (predecessor works too apparently.) two subcases:
        #    (a) node.right == S. this implies S.left is None (else S.left would be the
        #        successor), so S can replace node, adopting node's left child and
        #        keeping S's right child unchanged
        #
        #    (b) S is deeper in node.right's subtree.
        #
        #           we know this because the successor is either node.right.minimum() if
        #           node.right exists, or a traversal upwards, but node.right *must*
        #           exist because we have 2 children in this case. so S must be in
        #           node.right's subtree.
        #
        #        again we know S.left is None. we:
        #           - replace S with S.right
        #           - S replaces node
        s = node.successor()
        assert s is not None  # else node would be a leaf
        if node.right is s:
            # (a) case
            s.left = node.left
            _replace(node, s)
        else:
            # (b) case
            _replace(s, s.right)
            _replace(node, s)
            s.left = node.left
            s.right = node.right
        return self if node is not self else s


def _replace(node: BinaryTreeNode, replacement: Optional[BinaryTreeNode]):
    """replace node by setting node.parent's child to be replacement instead of node

    note that this doesn't alter any other children; that full operation is in delete().
    this just handles the fact that node.parent.child isn't defined because there's two
    children.
    """
    if node.parent is None:
        return  # nothing to wire up
    if node.parent.left is node:
        node.parent.left = replacement
    else:
        node.parent.right = replacement


def _traverse(subtree: Optional[BinaryTreeNode]):
    """Returns list of sorted values at and below subtree"""
    if subtree is None:
        return []

    return [
        *_traverse(subtree.left),
        subtree.value,
        *_traverse(subtree.right),
    ]


def _search(
    subtree: Optional[BinaryTreeNode], want: int
) -> Optional[BinaryTreeNode]:
    """Returns want's node if it's in subtree, else None"""
    if subtree is None:
        return None
    if subtree.value == want:
        return subtree
    if want < subtree.value:
        return _search(subtree.left, want)
    else:
        return _search(subtree.right, want)


class BST:
    """Somewhat awkward, likely unnecessary helper class to handle hanging onto the root
    reference and checking if it's None everywhere."""

    def __init__(self):
        self.root = None

    def insert(self, value: int):
        if self.root is None:
            self.root = BinaryTreeNode(value, None, None, None)
        else:
            self.root.insert(BinaryTreeNode(value, None, None, None))

    def traverse(self):
        return _traverse(self.root)

    def search(self, want: int):
        return _search(self.root, want)

    def minimum(self) -> Optional[BinaryTreeNode]:
        if self.root is None:
            return None
        return self.root.minimum()

    def maximum(self) -> Optional[BinaryTreeNode]:
        if self.root is None:
            return None
        return self.root.maximum()

    def delete(self, node: Optional[BinaryTreeNode]):
        if self.root is None:
            print("Can't delete node when tree root is None")
            return
        if node is None:
            return  # allow passing None just for type checking simplicity
        # NOTE: Could verify that node is in tree.
        self.root = self.root.delete(node)


def main() -> None:
    # shuffled [1, 10]
    vs = list(range(1, 11))
    random.shuffle(vs)
    print("original:     ", vs)

    # build into BST
    tree = BST()
    [tree.insert(v) for v in vs]

    # traverse
    print("BST traversal:", tree.traverse())
    print()

    # search
    print("BST search")
    print("1 in tree:    ", tree.search(1))
    print("7 in tree:    ", tree.search(7))
    print("10 in tree:   ", tree.search(10))
    print("0 in tree:    ", tree.search(0))
    print("-5 in tree:   ", tree.search(-5))
    print("53 in tree:   ", tree.search(53))
    print()

    print("BST Min/max")
    print("minimum:", tree.minimum())
    print("maximum:", tree.maximum())
    print()

    print("BST pred")
    print("5's predecessor:", tree.search(5).predecessor())  # type: ignore
    print("9's predecessor:", tree.search(9).predecessor())  # type: ignore
    print("1's predecessor:", tree.search(1).predecessor())  # type: ignore
    print("10's predecessor:", tree.search(10).predecessor())  # type: ignore
    print()

    print("BST suc")
    print("5's successor:", tree.search(5).successor())  # type: ignore
    print("9's successor:", tree.search(9).successor())  # type: ignore
    print("1's successor:", tree.search(1).successor())  # type: ignore
    print("10's successor:", tree.search(10).successor())  # type: ignore
    print()

    print("BST delete")
    print("root:", tree.root)
    tree.delete(tree.root)
    print("new root:", tree.root)
    print("traversal:", tree.traverse())


if __name__ == "__main__":
    main()

From my source repo.

"""Trie - a basic prefix tree.

Stores a value for each string (default just 1). This conveniently doubles as
indicating that we've reached a terminus (i.e., we've hit a stored word),
necessary for non-leaf termination (e.g., if storing all of 'i', 'in', and
'inn')
"""

# import code
from typing import Optional


class TrieNode:
    __slots__ = ("children", "value")
    """nice trick: for type checker to verify attrs"""

    def __init__(self):
        self.children: dict[str, TrieNode] = {}
        self.value: Optional[int] = None

    def insert(self, s: str, v: int = 1):
        """We store v (int) as payload data associated with s. For only storing
        a set of strings, it might as well be a single bit to indicate possible
        non-leaf termination (e.g., if storing all of 'i', 'in', and 'inn')."""
        # terminal
        if len(s) == 0:
            self.value = v
            return

        # any continuations
        c, rest = s[0], s[1:]
        if c not in self.children:
            self.children[c] = TrieNode()
        self.children[c].insert(rest, v)

    def traverse(self, accum: list[str] = []) -> list[str]:
        """Returns a list of all stored words."""
        res = []

        # terminal
        if self.value is not None:
            res.append("".join(accum) + f" ({self.value})")

        # any continuations
        for c, node in self.children.items():
            res.extend(node.traverse(accum + [c]))

        return res

    def _depth_sets(self, level=0):
        """helper for fun viz"""
        res = [(level, set(c for c in self.children.keys()))]
        for node in self.children.values():
            res.extend(node._depth_sets(level + 1))
        return res

    def print_depths(self):
        """fun viz (accumulate characters at each tree depth and print).

        In lieu of figuring out how to actually print trees "graphically" in text.
        """
        ss = self._depth_sets()
        for i in range(max(d for d, _cs in ss)):
            d_cs = set()
            for cs in [cs for d, cs in ss if d == i]:
                d_cs.update(cs)
            print(i, d_cs)

    def __repr__(self):
        return "\n".join(self.traverse())

    def search(self, s: str) -> Optional["TrieNode"]:
        """Returns node if it's found in the trie, else None."""
        if len(s) == 0:
            if self.value is not None:
                return self  # found
            return None  # not here (no terminal)
        c, rest = s[0], s[1:]
        if c not in self.children:
            return None  # no path
        return self.children[c].search(rest)  # continue

    def delete(self, s: str) -> bool:
        """returns whether this node can be fully deleted"""
        if len(s) == 0:
            # end of the line. mark no longer terminal.
            self.value = None

            # now, check whether there's children. if so, we leave everything
            # because continuations below us depend on us and our parents. but
            # if not, we can start propagating deletes upwards.
            return len(self.children) == 0

        # send down to children to delete
        c, rest = s[0], s[1:]
        if c not in self.children:
            raise ValueError(f"continuation {c} not found where expected")
        child_can_be_deleted = self.children[c].delete(rest)
        if child_can_be_deleted:
            del self.children[c]
            # continue propagating deletes up if we've no children and nonterminal
            return len(self.children) == 0 and self.value is None
        return False


def main() -> None:
    t = TrieNode()
    t.insert("hello", 47)
    t.insert("goodbye", 19)
    t.insert("i")
    print(t)
    t.print_depths()
    print()

    print("Search:")
    print("hello:", t.search("hello"))
    print("hell:", t.search("hell"))
    print("goodbye:", t.search("goodbye"))
    print("goodbyee:", t.search("goodbyee"))
    print("", t.search(""))
    print()

    print("Adding hell, good")
    t.insert("hell")
    t.insert("good")
    print(t)
    t.print_depths()
    print()

    print("Delete")
    print("delete hello")
    t.delete("hello")
    print(t)
    t.print_depths()
    print("delete good")
    t.delete("good")
    print(t)
    t.print_depths()
    print("delete goodbye")
    t.delete("goodbye")
    print(t)
    t.print_depths()

    # code.interact(local=dict(globals(), **locals()))


if __name__ == "__main__":
    main()

From my source repo.

I made this a little more complex because I wanted to implement both the O(rk) and O(k) approaches as subclasses. It’s a minor difference. The main algorithms are in search(), insert(), and delete() of the base class.

In implementing this, I didn’t reference any written material, only drew pictures and considered the cases. I couldn’t find a good reference description or implementation. So there certainly may be bugs. But my simple tests have worked correctly.

Before the full implementation, I’d like to pull out the central algorithms of search(), insert(), and delete(), and present them without comments or asserts or correctness checks so that we may gaze upon the core idea of the algorithm.

# Important: this is all pseudocode, because it's taken out of context.
# Imagine functions in a node class that holds:
#  - .terminal (i.e., does a word end here?)
#  - .children (some way of storing edges to children)

def _find_insertion_node(self, word: str) -> tuple[CTrieNode, str]:
    """returns node where word would be inserted, and how much is left."""
    if len(word) == 0:
        return (self, word)

    match = self._matches_edge(word)  # word substr exactly matching edge
    if match is None:
        return (self, word)
    substr, node = match
    remainder = word[len(substr) :]
    return node._find_insertion_node(remainder)

def search(self, word: str) -> bool:
    """Returns whether word is found in this subtree"""
    node, remainder = self._find_insertion_node(word)
    return len(remainder) == 0 and node.terminal

def insert(self, word: str):
    """inserts word into this subtree"""
    node, remainder = self._find_insertion_node(word)
    if len(remainder) == 0:
        node.terminal = True  # already have node at terminal
        return
    subnode, edge, common = node._overlaps_edge(remainder)
    if subnode is None:
        node._add_child(remainder, True)  # no match; fresh child
    else:
        # overlap, so split: (a) end at intermediate, or (b) new past
        intermediate_is_terminal = common == len(remainder)
        node._remove_edge(edge)
        intermediate = node._add_child(
            remainder[:common],
            intermediate_is_terminal,  # (a)
        )
        intermediate._add_child_node(edge[common:], subnode)
        if not intermediate_is_terminal:  # (b)
            intermediate._add_child(remainder[common:], True)

def delete(self, word: str):
    """deletes word"""
    parent, edge, node = self._find_parent(word)
    node.terminal = False
    if node._n_children() >= 2:
        pass  # keep node
    elif node._n_children() == 1:
        # link parent --> child (remove node)
        child_edge, child = node._get_only_child()
        parent._replace_edge(edge, edge + child_edge, child)
    elif node._n_children() == 0:
        parent._remove_edge(edge)  # remove node
        if parent._n_children() == 1:
            # link grandparent --> sibling (remove parent)
            grandparent, parent_edge, _ = self._find_parent(
                word[: -len(edge)]
            )
            sibling_edge, sibling = parent._get_only_child()
            grandparent._replace_edge(
                parent_edge, parent_edge + sibling_edge, sibling
            )

Core algorithms search(), insert(), and delete() for a compressed trie. See below for the full implementation.

Now, here’s the full implementation:

"""Compressed trie, AKA: compact prefix tree, Patricia tree/trie, radix tree/trie (ish).

This is a simple variant of a vanilla (character-per-node) trie where each node that is
an only child is merged with its parent.
"""

from abc import ABC, abstractmethod
from typing import Optional


def longest_common_start(s1: str, s2: str) -> int:
    """returns number of characters of longest common starting substring of s1 and s2.

    for example:
        0 1 2 3 4 5 6 7
        a r t e r i e s
        a r t
        -> returns 3

        0 1 2 3 4 5 6 7
        a r t e r i e s
        a r t i f a c t
        -> returns 3

        0 1 2 3 4 5
        a r t
        b a n a n a
        -> returns 0
    """
    min_len = min(len(s1), len(s2))
    for i in range(min_len):
        if s1[i] != s2[i]:
            return i
    return min_len


class CTrieNode(ABC):
    __slots__ = ("terminal",)

    def __init__(self, terminal: bool):
        self.terminal: bool = terminal

    # The following abstract methods allow a concrete node subclass to store
    # its edges and children in different data structures. Done this way
    # demonstrate the O(rk) and O(k) lookup approaches.

    @abstractmethod
    def _matches_edge(self, word: str) -> Optional[tuple[str, "CTrieNode"]]:
        """returns (substring, CTrieNode) if there's a substring of word that
        exactly matches an edge, else None.
        """
        pass

    @abstractmethod
    def _overlaps_edge(
        self, word: str
    ) -> tuple[Optional["CTrieNode"], str, int]:
        """Returns any overlapping edge's (node, edge, overlap len),
        or (None, "", 0) if there's none.

        At most 1 edge will have any common start overlap with word b/c no edges
        have any prefix overlap because that's how a prefix data structure works.
        """
        pass

    @abstractmethod
    def _add_child(self, edge: str, terminal: bool) -> "CTrieNode":
        """adds edge pointing to a new node w/ specified terminal set. returns child.

        may validate no shared prefixes, so if modifying, delete before calling this.
        """
        pass

    @abstractmethod
    def _add_child_node(self, edge: str, child: "CTrieNode") -> "CTrieNode":
        """adds new edge pointing to child. returns child.

        may validate no shared prefixes, so if modifying, delete before calling this.
        """
        pass

    @abstractmethod
    def _remove_edge(self, edge: str) -> None:
        """removes edge"""
        pass

    @abstractmethod
    def all_children(self) -> list[tuple[str, "CTrieNode"]]:
        pass

    @abstractmethod
    def _depth_sets(self, level=0) -> list[tuple[int, set[str]]]:
        """helper for fun viz. returns list of [depth, set(strs at depth) ]"""
        pass

    @abstractmethod
    def _n_children(self) -> int:
        """conceptually, self.terminal counts as another child, so its included in the
        total."""
        pass

    @abstractmethod
    def _get_only_child(self) -> tuple[str, "CTrieNode"]:
        """gets the only child's (edge, child). should only be called when
        self._n_children() == 1, and not self.terminal. only used before deeply
        modifying."""
        pass

    # the following happen to be the same implementation, but could be broken
    # out into subclasses if needed.

    def _replace_edge(
        self,
        old_edge: str,
        new_edge: str,
        node: "CTrieNode",
    ) -> None:
        """replaces old_edge with new_edge pointing to node"""
        self._remove_edge(old_edge)
        self._add_child_node(new_edge, node)

    # the following are the main concrete operations implementing the data structure.

    def _find_insertion_node(self, word: str) -> tuple["CTrieNode", str]:
        """returns node where word would be inserted, and how much is left

        returns (node, remainder) where remainder is any trailing substr of word. Note
        that if the word was found, node.terminal() == True and remainder == ''.
        """
        # terminal case
        if len(word) == 0:
            return (self, word)

        # search edges case
        match = self._matches_edge(word)
        if match is None:
            return (self, word)
        substr, node = match
        remainder = word[len(substr) :]
        return node._find_insertion_node(remainder)

    def search(self, word: str) -> bool:
        """Returns whether word is found in this subtree"""
        node, remainder = self._find_insertion_node(word)
        return len(remainder) == 0 and node.terminal

    def insert(self, word: str):
        """inserts word into this subtree"""
        node, remainder = self._find_insertion_node(word)
        if len(remainder) == 0:
            # already have node at terminal! ensure terminal True (already is if exists)
            node.terminal = True
            return

        # otherwise, there's some remainder left. two main cases: (1) there's
        # no matching substr below, so we make a fresh child. (2) we have a
        # partial match with an edge, so we'll have to split.
        subnode, edge, lcs = node._overlaps_edge(remainder)
        if subnode is None:
            # (1). means lcs == 0, so we add all remainder
            assert len(edge) == 0 and lcs == 0
            node._add_child(remainder, True)
        else:
            # (2) split. two suboptions:
            #     (a) our intermediate node is terminal if lcs == len(remainder).
            #         (like []--"artifice"-->[] + "art" = []--"art"-->[]--"ifice"-->[])
            #
            #     (b) our intermediate node isn't terminal and we need another for rest.
            # add intermediate node
            intermediate_is_terminal = lcs == len(remainder)
            node._remove_edge(edge)  # do now for sanity checks
            intermediate = node._add_child(
                remainder[:lcs],
                intermediate_is_terminal,  # for (a)
            )
            # always add back the existing subnode with a shortened edge
            intermediate._add_child_node(edge[lcs:], subnode)
            # for (b), add a new node for this word
            if not intermediate_is_terminal:
                intermediate._add_child(remainder[lcs:], True)

    def _find_parent(self, word: str) -> tuple["CTrieNode", str, "CTrieNode"]:
        """returns (parent, edge, node), where node specifies word. helper for delete().

        very similar to _find_insertion_node(), but we don't have parent
        references. so if we're to delete a node, we need to be looking down
        towards it.
        """
        # search edges case
        match = self._matches_edge(word)
        if match is None:
            raise ValueError(f"word (remainder) '{word}' not found in ctrie")
        substr, node = match
        remainder = word[len(substr) :]
        if len(remainder) == 0:
            return self, substr, node  # found case
        return node._find_parent(remainder)  # recurse case

    def delete(self, word: str):
        print(f"deleting '{word}'")
        parent, edge, node = self._find_parent(word)
        if node.terminal is False:
            raise ValueError(f"Can't delete '{word}' as it wasn't stored")
        node.terminal = False
        if node._n_children() >= 2:
            # easy case: node still used by children. done.
            print("  delete: 2 children case (easy)")
            pass
        elif node._n_children() == 1:
            # node was a terminal with a single option below (conceptually, 2
            # -> 1 children). now that it's no longer a terminal, its should
            # be compressed with its parent. example:
            #
            #             deleted "art", no longer Terminal
            #             v
            # []--"art"-->[]--"ifice"-->[] (-->)
            # parent      node          child
            #
            # should now become:
            #
            # []--"artifice"-->[]
            # parent          child
            #
            # we don't need to worry about parent, because parent must have started with
            # >= 2 children (including 1 if it's terminal). node was 1 of its children,
            # and now child will *still* be 1 of its children.
            print("  delete: 1 child case (medium)")
            child_edge, child = node._get_only_child()
            parent._replace_edge(edge, edge + child_edge, child)
        elif node._n_children() == 0:
            # node had no children, so it can be removed. this means parent's
            # number of children decreases, which means it might need to be
            # merged with its parent if it has 1 left (node's sibling). this will
            # make grandparent point directly to sibling.
            print("  delete: 0 children case (hard)")
            parent._remove_edge(edge)
            if parent._n_children() == 1:
                # example: deleting "artifice":
                #
                #                 +----"ful"----> [] sibling
                #                 |
                # [] ---"art"---> [] --"ifice"--> [] node
                # grandparent   parent
                #
                #                 +----"ful"----> [] sibling
                #                 |
                # [] ---"art"---> []
                # grandparent   parent <-- ! only 1 child! merge w/ grandparent
                #
                # [] ---"artful"---> []
                # grandparent      sibling

                # NOTE: this is inefficient: finding grandparent! should put
                # in _find_parent() but it's annoying to track grandparents.
                grandparent, parent_edge, parent_ = self._find_parent(
                    word[: -len(edge)]
                )
                assert parent_ is parent
                sibling_edge, sibling = parent._get_only_child()
                grandparent._replace_edge(
                    parent_edge, parent_edge + sibling_edge, sibling
                )

        else:
            assert False  # should never have < 0 children

    def traverse(self, prefix: list[str] = []) -> list[str]:
        """O(n), return a list of all stored strings"""
        res = []
        if self.terminal:
            res.append("".join(prefix))
        for edge, child in self.all_children():
            res.extend(child.traverse(prefix + [edge]))
        return res

    def __repr__(self) -> str:
        """O(n), returns all stored strings newline-separated"""
        return "\n".join(self.traverse())

    def print_depths(self) -> None:
        """fun viz (accumulate characters at each tree depth and print).

        In lieu of figuring out how to actually print trees "graphically" in text.
        """
        ss = self._depth_sets()
        for i in range(max(d for d, _cs in ss)):
            d_cs = set()
            for cs in [cs for d, cs in ss if d == i]:
                d_cs.update(cs)
            print(i, d_cs)


class CTrieNodeEdgeIterate(CTrieNode):
    """This is a data structure for a Compressed Trie Node that looks up edges by
    iterating over them, which is O(rk), where r is the radix, AKA the set of unique
    characters, and k is the length of a key (stored word). (For example, if storing
    lowercase English letters only, r = 26, and when storing 'apple', k = 5)

    This class contains the methods for getting/storing edges, and doing substring
    lookups. The core data structure methods search(), insert(), and delete() are
    implemented on the superclass CTrieNode.
    """

    __slots__ = ("_children",)

    def __init__(self, terminal: bool):
        super().__init__(terminal)
        self._children: dict[str, CTrieNode] = {}

    def _matches_edge(self, word: str) -> Optional[tuple[str, "CTrieNode"]]:
        """
        This is one of the core O(rk) operations we can speed up to O(k)
        """
        for substr, node in self._children.items():  # O(r)
            if word.startswith(substr):  # O(k)
                return (substr, node)
        return None

    def _overlaps_edge(
        self, word: str
    ) -> tuple[Optional["CTrieNode"], str, int]:
        """
        This is one of the core O(rk) operations we can speed up to O(k)
        """
        # due to that this is still a prefix trie, only one child will have a
        # common start > 0. because if multiple children did, that shared
        # prefix should already have been split down into its own node.
        for edge, node in self._children.items():  # O(r)
            lcs = longest_common_start(edge, word)  # O(k)
            if lcs > 0:
                return node, edge, lcs
        return None, "", 0

    def _add_child(self, edge: str, terminal: bool) -> "CTrieNode":
        return self._add_child_node(edge, CTrieNodeEdgeIterate(terminal))

    def _add_child_node(self, edge: str, child: "CTrieNode") -> "CTrieNode":
        self._children[edge] = child
        assert sum(edge[0] == w[0] for w in self._children.keys()) == 1
        return child

    def _remove_edge(self, edge: str) -> None:
        del self._children[edge]

    def all_children(self) -> list[tuple[str, "CTrieNode"]]:
        return list(self._children.items())

    def _depth_sets(self, level=0) -> list[tuple[int, set[str]]]:
        res = [
            (
                level,
                set(f"{c} ({n.terminal})" for c, n in self._children.items()),
            )
        ]
        for node in self._children.values():
            res.extend(node._depth_sets(level + 1))
        return res

    def _n_children(self) -> int:
        return len(self._children) + (1 if self.terminal else 0)

    def _get_only_child(self) -> tuple[str, "CTrieNode"]:
        assert not self.terminal and self._n_children() == 1
        for edge, node in self._children.items():
            return edge, node
        assert False


class CTrieNodeCharLookup(CTrieNode):
    """This is a data structure for a Compressed Trie Node that looks up edges by
    indexing by the first letter, which is O(1), then comparing strings, which is O(k).

    This class contains the methods for getting/storing edges, and doing substring
    lookups. The core data structure methods search(), insert(), and delete() are
    implemented on the superclass CTrieNode.
    """

    __slots__ = ("_children",)

    def __init__(self, terminal: bool):
        super().__init__(terminal)
        self._children: dict[str, tuple[str, CTrieNode]] = {}
        """dict from {edge's first char: (edge, node)}, like {'a': ('apple', node)}"""

    def _matches_edge(self, word: str) -> Optional[tuple[str, "CTrieNode"]]:
        """
        This is one of the core O(k) operations we sped up from O(rk)
        """
        if len(word) > 0 and word[0] in self._children:  # O(1)
            edge, node = self._children[word[0]]
            if word.startswith(edge):  # O(k)
                return (edge, node)
        return None

    def _overlaps_edge(
        self, word: str
    ) -> tuple[Optional["CTrieNode"], str, int]:
        """
        This is one of the core O(k) operations we sped up from O(rk)
        """
        if len(word) > 0 and word[0] in self._children:  # O(1)
            edge, node = self._children[word[0]]
            lcs = longest_common_start(edge, word)  # O(k)
            if lcs > 0:
                return node, edge, lcs
        return None, "", 0

    def _add_child(self, edge: str, terminal: bool) -> "CTrieNode":
        return self._add_child_node(edge, CTrieNodeCharLookup(terminal))

    def _add_child_node(self, edge: str, child: "CTrieNode") -> "CTrieNode":
        first = edge[0]
        if first in self._children:
            raise ValueError(
                f"cannot add edge {edge} when {self._children[first][0]} already stored"
            )
        self._children[first] = (edge, child)
        return child

    def _remove_edge(self, edge: str) -> None:
        del self._children[edge[0]]

    def all_children(self) -> list[tuple[str, "CTrieNode"]]:
        return list(self._children.values())

    def _depth_sets(self, level=0) -> list[tuple[int, set[str]]]:
        res = [
            (
                level,
                set(
                    f"{c} ({n.terminal})" for c, n in self._children.values()
                ),
            )
        ]
        for _, node in self._children.values():
            res.extend(node._depth_sets(level + 1))
        return res

    def _n_children(self) -> int:
        return len(self._children) + (1 if self.terminal else 0)

    def _get_only_child(self) -> tuple[str, "CTrieNode"]:
        assert not self.terminal and self._n_children() == 1
        for edge, node in self._children.values():
            return edge, node
        assert False


def main() -> None:
    # test helpers
    assert longest_common_start("art", "arteries") == 3
    assert longest_common_start("artifact", "arteries") == 3
    assert longest_common_start("art", "banana") == 0

    # to switch between node implementations
    # Node = CTrieNodeEdgeIterate
    Node = CTrieNodeCharLookup

    t1 = Node(False)
    t1.insert("artifact")
    t1.insert("artifice")
    t1.insert("art")
    t1.insert("artificial")
    print(t1)
    t1.print_depths()

    # ctrie.delete("art")  # 1 child (medium)
    # ctrie.delete("artificial") # 0 children (hard)
    # 2 children (following two lines): easy
    t1.insert("artful")
    t1.delete("art")
    print(t1)
    t1.print_depths()

    t2 = Node(False)
    t2.insert("julie")
    t2.insert("july")
    # t.delete("jul")  # value error
    t2.insert("julius")
    t2.delete("julie")
    # t.insert("abcdefghijklmnopqrstuvwxyz")
    # t.insert("abra")
    t2.print_depths()


if __name__ == "__main__":
    main()

From my source repo.