mirror of
https://github.com/TheAlgorithms/Python.git
synced 2025-01-18 16:27:02 +00:00
parent
02c0daf9e5
commit
94380a17a8
129
data_structures/binary tree/treap.py
Normal file
129
data_structures/binary tree/treap.py
Normal file
|
@ -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()
|
Loading…
Reference in New Issue
Block a user