Python/data_structures/binary tree/treap.py

from random import random
from typing import Tuple


class Node:
    """
    Treap's node
    Treap is a binary tree by key and heap by priority
    """
    def __init__(self, key: int):
        self.key = key
        self.prior = random()
        self.l = None
        self.r = None


def split(root: Node, key: int) -> Tuple[Node, Node]:
    """
    We split current tree into 2 trees with key:

    Left tree contains all keys less than split key.
    Right tree contains all keys greater or equal, than split key
    """
    if root is None:  # None tree is split into 2 Nones
        return (None, None)
    if root.key >= key:
        """
        Right tree's root will be current node.
        Now we split(with the same key) current node's left son
        Left tree: left part of that split
        Right tree's left son: right part of that split
        """
        l, root.l = split(root.l, key)
        return (l, root)
    else:
        """
        Just symmetric to previous case
        """
        root.r, r = split(root.r, key)
        return (root, r)


def merge(left: Node, right: Node) -> Node:
    """
    We merge 2 trees into one.
    Note: all left tree's keys must be less than all right tree's
    """
    if (not left) or (not right):
        """
        If one node is None, return the other
        """
        return left or right
    if left.key > right.key:
        """
        Left will be root because it has more priority
        Now we need to merge left's right son and right tree
        """
        left.r = merge(left.r, right)
        return left
    else:
        """
        Symmetric as well
        """
        right.l = merge(left, right.l)
        return right


def insert(root: Node, key: int) -> Node:
    """
    Insert element

    Split current tree with a key into l, r,
    Insert new node into the middle
    Merge l, node, r into root
    """
    node = Node(key)
    l, r = split(root, key)
    root = merge(l, node)
    root = merge(root, r)
    return root


def erase(root: Node, key: int) -> Node:
    """
    Erase element

    Split all nodes with keys less into l,
    Split all nodes with keys greater into r.
    Merge l, r
    """
    l, r = split(root, key)
    _, r = split(r, key + 1)
    return merge(l, r)


def node_print(root: Node):
    """
    Just recursive print of a tree
    """
    if not root:
        return
    node_print(root.l)
    print(root.key, end=" ")
    node_print(root.r)


def interactTreap():
    """
    Commands:
    + key to add key into treap
    - key to erase all nodes with key

    After each command, program prints treap
    """
    root = None
    while True:
        cmd = input().split()
        cmd[1] = int(cmd[1])
        if cmd[0] == "+":
            root = insert(root, cmd[1])
        elif cmd[0] == "-":
            root = erase(root, cmd[1])
        else:
            print("Unknown command")
        node_print(root)


if __name__ == "__main__":
    interactTreap()