forked from aosabook/500lines
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathbinary_tree.py
More file actions
142 lines (124 loc) · 4.51 KB
/
Copy pathbinary_tree.py
File metadata and controls
142 lines (124 loc) · 4.51 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
import pickle
from dbdb.logical import LogicalBase, ValueRef
class BinaryNode(object):
@classmethod
def from_node(cls, node, **kwargs):
length = node.length
if 'left_ref' in kwargs:
length += kwargs['left_ref'].length - node.left_ref.length
if 'right_ref' in kwargs:
length += kwargs['right_ref'].length - node.right_ref.length
return cls(
left_ref=kwargs.get('left_ref', node.left_ref),
key=kwargs.get('key', node.key),
value_ref=kwargs.get('value_ref', node.value_ref),
right_ref=kwargs.get('right_ref', node.right_ref),
length=length,
)
def __init__(self, left_ref, key, value_ref, right_ref, length):
self.left_ref = left_ref
self.key = key
self.value_ref = value_ref
self.right_ref = right_ref
self.length = length
def store_refs(self, storage):
self.value_ref.store(storage)
self.left_ref.store(storage)
self.right_ref.store(storage)
class BinaryNodeRef(ValueRef):
def prepare_to_store(self, storage):
if self._referent:
self._referent.store_refs(storage)
@property
def length(self):
if self._referent is None and self._address:
raise RuntimeError('Asking for BinaryNodeRef length of unloaded node')
if self._referent:
return self._referent.length
else:
return 0
@staticmethod
def referent_to_string(referent):
return pickle.dumps({
'left': referent.left_ref.address,
'key': referent.key,
'value': referent.value_ref.address,
'right': referent.right_ref.address,
'length': referent.length,
})
@staticmethod
def string_to_referent(string):
d = pickle.loads(string)
return BinaryNode(
BinaryNodeRef(address=d['left']),
d['key'],
ValueRef(address=d['value']),
BinaryNodeRef(address=d['right']),
d['length'],
)
class BinaryTree(LogicalBase):
node_ref_class = BinaryNodeRef
def _get(self, node, key):
while node is not None:
if key < node.key:
node = self._follow(node.left_ref)
elif node.key < key:
node = self._follow(node.right_ref)
else:
return self._follow(node.value_ref)
raise KeyError
def _insert(self, node, key, value_ref):
if node is None:
new_node = BinaryNode(
self.node_ref_class(), key, value_ref, self.node_ref_class(), 1)
elif key < node.key:
new_node = BinaryNode.from_node(
node,
left_ref=self._insert(
self._follow(node.left_ref), key, value_ref))
elif node.key < key:
new_node = BinaryNode.from_node(
node,
right_ref=self._insert(
self._follow(node.right_ref), key, value_ref))
else:
new_node = BinaryNode.from_node(node, value_ref=value_ref)
return self.node_ref_class(referent=new_node)
def _delete(self, node, key):
if node is None:
raise KeyError
elif key < node.key:
new_node = BinaryNode.from_node(
node,
left_ref=self._delete(
self._follow(node.left_ref), key))
elif node.key < key:
new_node = BinaryNode.from_node(
node,
right_ref=self._delete(
self._follow(node.right_ref), key))
else:
left = self._follow(node.left_ref)
right = self._follow(node.right_ref)
if left and right:
replacement = self._find_max(left)
left_ref = self._delete(
self._follow(node.left_ref), replacement.key)
new_node = BinaryNode(
left_ref,
replacement.key,
replacement.value_ref,
node.right_ref,
left_ref.length + node.right_ref.length + 1,
)
elif left:
return node.left_ref
else:
return node.right_ref
return self.node_ref_class(referent=new_node)
def _find_max(self, node):
while True:
next_node = self._follow(node.right_ref)
if next_node is None:
return node
node = next_node