blog/content/posts/2024-07-27-treap-revisited/index.md

title	date	draft	description	tags	categories	series	favorite	disable_feed
Treap, revisited	2024-07-27T14:12:27+01:00	false	An even simpler BST	algorithms data structures python	programming	Cool algorithms	false	false

My [last post]({{< relref "../2024-07-20-treap/index.md" >}}) about the Treap showed an implementation using tree rotations, as is commonly done with AVL Trees and Red Black Trees.

But the Treap lends itself well to a simple and elegant implementation with no tree rotations. This makes it especially easy to implement the removal of a key, rather than the fiddly process of deletion using tree rotations.

Implementation

All operations on the tree will be implemented in terms of two fundamental operations: split and merge.

We'll be reusing the same structures as in the last post, so let's skip straight to implementing those fundaments, and building on them for insert and delete.

Split

Splitting a tree means taking a key, and getting the following output:

• a left node, root of the tree of all keys lower than the input.
• an extracted node which corresponds to the input key.
• a right node, root of the tree of all keys higher than the input.

type OptionalNode[K, V] = Node[K, V] | None

class SplitResult(NamedTuple):
    left: OptionalNode
    node: OptionalNode
    right: OptionalNode

def split(root: OptionalNode[K, V], key: K) -> SplitResult:
    # Base case, empty tree
    if root is None:
        return SplitResult(None, None, None)
    # If we found the key, simply extract left and right
    if root.key == key:
        left, right = root.left, root.right
        root.left, root.right = None, None
        return SplitResult(left, root, right)
    # Otherwise, recurse on the corresponding side of the tree
    if root.key < key:
        left, node, right = split(root.right, key)
        root.right = left
        return SplitResult(root, node, right)
    if key < root.key:
        left, node, right = split(root.left, key)
        root.left = right
        return SplitResult(left, node, root)
    raise RuntimeError("Unreachable")

Merge

Merging a left and right tree means (cheaply) building a new tree containing both of them. A pre-condition for merging is that the left tree is composed entirely of nodes that are lower than any key in right (i.e: as in left and right after a split).

def merge(
    left: OptionalNode[K, V],
    right: OptionalNode[K, V],
) -> OptionalNode[K, V]:
    # Base cases, left or right being empty
    if left is None:
        return right
    if right is None:
        return left
    # Left has higher priority, it must become the root node
    if left.priority >= right.priority:
        # We recursively reconstruct its right sub-tree
        left.right = merge(left.right, right)
        return left
    # Right has higher priority, it must become the root node
    if left.priority < right.priority:
        # We recursively reconstruct its left sub-tree
        right.left = merge(left, right.left)
        return right
    raise RuntimeError("Unreachable")

Insertion

Inserting a node into the tree is done in two steps:

1. split the tree to isolate the middle insertion point
2. merge it back up to form a full tree with the inserted key

def insert(self, key: K, value: V) -> bool:
    # `left` and `right` come before/after the key
    left, node, right = split(self._root, key)
    was_updated: bool
    # Create the node, or update its value, if the key was already in the tree
    if node is None:
        node = Node(key, value)
        was_updated = False
    else:
        node.value = value
        was_updated = True
    # Rebuild the tree with a couple of merge operations
    self._root = merge(left, merge(node, right))
    # Signal whether the key was already in the key
    return was_updated

Removal

Removing a key from the tree is similar to inserting a new key, and forgetting to insert it back: simply split the tree and merge it back without the extracted middle node.

def remove(self, key: K) -> bool:
    # `node` contains the key, or `None` if the key wasn't in the tree
    left, node, right = split(self._root, key)
    # Put the tree back together, without the extract node
    self._root = merge(left, right)
    # Signal whether `key` was mapped in the tree
    return node is not None