Added treap

data_structures/binary tree/treap.py

@@ -0,0 +1,129 @@
+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()