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

3.2 KiB

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

def split(
    root: OptionalNode[K, V],
    key: K,
) -> tuple[OptionalNode[K, V], OptionalNode[K, V], OptionalNode[K, V]]:
    # Base case, empty tree
    if root is None:
        return 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 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 root, node, right
    if key < root.key:
        left, node, right = split(root.left, key)
        root.left = right
        return 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")