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

2.3 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

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")