Skip to content

Commit 94380a1

Browse files
Flygrounderpoyea
authored andcommitted
Added treap (TheAlgorithms#797)
* Added treap * Added comments to treap
1 parent 02c0daf commit 94380a1

File tree

1 file changed

+129
-0
lines changed

1 file changed

+129
-0
lines changed

data_structures/binary tree/treap.py

+129
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
from random import random
2+
from typing import Tuple
3+
4+
5+
class Node:
6+
"""
7+
Treap's node
8+
Treap is a binary tree by key and heap by priority
9+
"""
10+
def __init__(self, key: int):
11+
self.key = key
12+
self.prior = random()
13+
self.l = None
14+
self.r = None
15+
16+
17+
def split(root: Node, key: int) -> Tuple[Node, Node]:
18+
"""
19+
We split current tree into 2 trees with key:
20+
21+
Left tree contains all keys less than split key.
22+
Right tree contains all keys greater or equal, than split key
23+
"""
24+
if root is None: # None tree is split into 2 Nones
25+
return (None, None)
26+
if root.key >= key:
27+
"""
28+
Right tree's root will be current node.
29+
Now we split(with the same key) current node's left son
30+
Left tree: left part of that split
31+
Right tree's left son: right part of that split
32+
"""
33+
l, root.l = split(root.l, key)
34+
return (l, root)
35+
else:
36+
"""
37+
Just symmetric to previous case
38+
"""
39+
root.r, r = split(root.r, key)
40+
return (root, r)
41+
42+
43+
def merge(left: Node, right: Node) -> Node:
44+
"""
45+
We merge 2 trees into one.
46+
Note: all left tree's keys must be less than all right tree's
47+
"""
48+
if (not left) or (not right):
49+
"""
50+
If one node is None, return the other
51+
"""
52+
return left or right
53+
if left.key > right.key:
54+
"""
55+
Left will be root because it has more priority
56+
Now we need to merge left's right son and right tree
57+
"""
58+
left.r = merge(left.r, right)
59+
return left
60+
else:
61+
"""
62+
Symmetric as well
63+
"""
64+
right.l = merge(left, right.l)
65+
return right
66+
67+
68+
def insert(root: Node, key: int) -> Node:
69+
"""
70+
Insert element
71+
72+
Split current tree with a key into l, r,
73+
Insert new node into the middle
74+
Merge l, node, r into root
75+
"""
76+
node = Node(key)
77+
l, r = split(root, key)
78+
root = merge(l, node)
79+
root = merge(root, r)
80+
return root
81+
82+
83+
def erase(root: Node, key: int) -> Node:
84+
"""
85+
Erase element
86+
87+
Split all nodes with keys less into l,
88+
Split all nodes with keys greater into r.
89+
Merge l, r
90+
"""
91+
l, r = split(root, key)
92+
_, r = split(r, key + 1)
93+
return merge(l, r)
94+
95+
96+
def node_print(root: Node):
97+
"""
98+
Just recursive print of a tree
99+
"""
100+
if not root:
101+
return
102+
node_print(root.l)
103+
print(root.key, end=" ")
104+
node_print(root.r)
105+
106+
107+
def interactTreap():
108+
"""
109+
Commands:
110+
+ key to add key into treap
111+
- key to erase all nodes with key
112+
113+
After each command, program prints treap
114+
"""
115+
root = None
116+
while True:
117+
cmd = input().split()
118+
cmd[1] = int(cmd[1])
119+
if cmd[0] == "+":
120+
root = insert(root, cmd[1])
121+
elif cmd[0] == "-":
122+
root = erase(root, cmd[1])
123+
else:
124+
print("Unknown command")
125+
node_print(root)
126+
127+
128+
if __name__ == "__main__":
129+
interactTreap()

0 commit comments

Comments
 (0)