forked from TheAlgorithms/Python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtreap.py
More file actions
182 lines (149 loc) · 4.74 KB
/
Copy pathtreap.py
File metadata and controls
182 lines (149 loc) · 4.74 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
from __future__ import annotations
from random import random
class Node:
"""
Treap's node
Treap is a binary tree by value and heap by priority
"""
def __init__(self, value: int | None = None):
self.value = value
self.prior = random()
self.left: Node | None = None
self.right: Node | None = None
def __repr__(self) -> str:
from pprint import pformat
if self.left is None and self.right is None:
return f"'{self.value}: {self.prior:.5}'"
else:
return pformat(
{f"{self.value}: {self.prior:.5}": (self.left, self.right)}, indent=1
)
def __str__(self) -> str:
value = str(self.value) + " "
left = str(self.left or "")
right = str(self.right or "")
return value + left + right
def split(root: Node | None, value: int) -> tuple[Node | None, Node | None]:
"""
We split current tree into 2 trees with value:
Left tree contains all values less than split value.
Right tree contains all values greater or equal, than split value
"""
if root is None: # None tree is split into 2 Nones
return None, None
elif root.value is None:
return None, None
else:
if value < root.value:
"""
Right tree's root will be current node.
Now we split(with the same value) current node's left son
Left tree: left part of that split
Right tree's left son: right part of that split
"""
left, root.left = split(root.left, value)
return left, root
else:
"""
Just symmetric to previous case
"""
root.right, right = split(root.right, value)
return root, right
def merge(left: Node | None, right: Node | None) -> Node | None:
"""
We merge 2 trees into one.
Note: all left tree's values 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
elif left.prior < right.prior:
"""
Left will be root because it has more priority
Now we need to merge left's right son and right tree
"""
left.right = merge(left.right, right)
return left
else:
"""
Symmetric as well
"""
right.left = merge(left, right.left)
return right
def insert(root: Node | None, value: int) -> Node | None:
"""
Insert element
Split current tree with a value into left, right,
Insert new node into the middle
Merge left, node, right into root
"""
node = Node(value)
left, right = split(root, value)
return merge(merge(left, node), right)
def erase(root: Node | None, value: int) -> Node | None:
"""
Erase element
Split all nodes with values less into left,
Split all nodes with values greater into right.
Merge left, right
"""
left, right = split(root, value - 1)
_, right = split(right, value)
return merge(left, right)
def inorder(root: Node | None) -> None:
"""
Just recursive print of a tree
"""
if not root: # None
return
else:
inorder(root.left)
print(root.value, end=",")
inorder(root.right)
def interact_treap(root: Node | None, args: str) -> Node | None:
"""
Commands:
+ value to add value into treap
- value to erase all nodes with value
>>> root = interact_treap(None, "+1")
>>> inorder(root)
1,
>>> root = interact_treap(root, "+3 +5 +17 +19 +2 +16 +4 +0")
>>> inorder(root)
0,1,2,3,4,5,16,17,19,
>>> root = interact_treap(root, "+4 +4 +4")
>>> inorder(root)
0,1,2,3,4,4,4,4,5,16,17,19,
>>> root = interact_treap(root, "-0")
>>> inorder(root)
1,2,3,4,4,4,4,5,16,17,19,
>>> root = interact_treap(root, "-4")
>>> inorder(root)
1,2,3,5,16,17,19,
>>> root = interact_treap(root, "=0")
Unknown command
"""
for arg in args.split():
if arg[0] == "+":
root = insert(root, int(arg[1:]))
elif arg[0] == "-":
root = erase(root, int(arg[1:]))
else:
print("Unknown command")
return root
def main() -> None:
"""After each command, program prints treap"""
root = None
print(
"enter numbers to create a tree, + value to add value into treap, "
"- value to erase all nodes with value. 'q' to quit. "
)
args = input()
while args != "q":
root = interact_treap(root, args)
print(root)
args = input()
print("good by!")
if __name__ == "__main__":
import doctest
doctest.testmod()
main()