FilesExpand file tree

master

/

treap.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

182 lines (149 loc) · 4.74 KB

master

/

treap.py

Copy path

Top

File metadata and controls

• Code
• Blame

182 lines (149 loc) · 4.74 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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()