1 | from abc import ABCMeta |
||
2 | from pyalgs.algorithms.commons.util import cmp |
||
3 | from pyalgs.data_structures.commons.queue import Queue |
||
4 | |||
5 | |||
6 | class Node(object): |
||
7 | key = None |
||
8 | value = None |
||
9 | |||
10 | left = None |
||
11 | right = None |
||
12 | |||
13 | count = 0 |
||
14 | red = 1 |
||
15 | |||
16 | def __init__(self, key, value, red=None): |
||
17 | self.key = key |
||
18 | self.value = value |
||
19 | self.count = 1 |
||
20 | |||
21 | if red is not None: |
||
22 | self.red = red |
||
23 | else: |
||
24 | self.red = 0 |
||
25 | |||
26 | |||
27 | def _count(x): |
||
28 | if x is None: |
||
29 | return 0 |
||
30 | return x.count |
||
31 | |||
32 | |||
33 | class BinarySearchTree(object): |
||
34 | __metaclass__ = ABCMeta |
||
35 | |||
36 | root = None |
||
37 | |||
38 | def put(self, key, value): |
||
39 | self.root = self._put(self.root, key, value) |
||
40 | |||
41 | def _put(self, x, key, value): |
||
42 | if x is None: |
||
43 | return Node(key, value) |
||
44 | |||
45 | compared = cmp(key, x.key) |
||
46 | |||
47 | if compared < 0: |
||
48 | x.left = self._put(x.left, key, value) |
||
49 | elif compared > 0: |
||
50 | x.right = self._put(x.right, key, value) |
||
51 | else: |
||
52 | x.value = value |
||
53 | |||
54 | x.count = 1 + _count(x.left) + _count(x.right) |
||
55 | |||
56 | return x |
||
57 | |||
58 | def get(self, key): |
||
59 | x = self._get(self.root, key) |
||
60 | if x is None: |
||
61 | return None |
||
62 | return x.value |
||
63 | |||
64 | def _get(self, x, key): |
||
65 | if x is None: |
||
66 | return None |
||
67 | compared = cmp(key, x.key) |
||
68 | if compared < 0: |
||
69 | return self._get(x.left, key) |
||
70 | elif compared > 0: |
||
71 | return self._get(x.right, key) |
||
72 | else: |
||
73 | return x |
||
74 | |||
75 | def size(self): |
||
76 | return _count(self.root) |
||
77 | |||
78 | def is_empty(self): |
||
79 | return self.root is None |
||
80 | |||
81 | def delete(self, key): |
||
82 | self.root = self._delete(self.root, key) |
||
83 | |||
84 | def _delete(self, x, key): |
||
85 | if x is None: |
||
86 | return None |
||
87 | |||
88 | compared = cmp(key, x.key) |
||
89 | View Code Duplication | if compared < 0: |
|
0 ignored issues
–
show
Duplication
introduced
by
Loading history...
|
|||
90 | x.left = self._delete(x.left, key) |
||
91 | elif compared > 0: |
||
92 | x.right = self._delete(x.right, key) |
||
93 | else: |
||
94 | if x.left is None: |
||
95 | return x.right |
||
96 | elif x.right is None: |
||
97 | return x.left |
||
98 | else: |
||
99 | m = self.min(x.right) |
||
100 | m.right = self.del_min(x.right) |
||
101 | m.left = x.left |
||
102 | |||
103 | x = m |
||
104 | |||
105 | x.count = 1 + _count(x.left) + _count(x.right) |
||
106 | return x |
||
107 | |||
108 | def min(self, x): |
||
109 | if x.left is None: |
||
110 | return x |
||
111 | return self.min(x.left) |
||
112 | |||
113 | def del_min(self, x): |
||
114 | if x.left is None: |
||
115 | return x.right |
||
116 | x.left = self.del_min(x.left) |
||
117 | return x |
||
118 | |||
119 | def contains_key(self, x): |
||
120 | return self.get(x) is not None |
||
121 | |||
122 | def keys(self): |
||
123 | queue = Queue.create() |
||
124 | self.collect_keys(self.root, queue) |
||
125 | return queue.iterate() |
||
126 | |||
127 | def collect_keys(self, x, queue): |
||
128 | if x is None: |
||
129 | return |
||
130 | |||
131 | self.collect_keys(x.left, queue) |
||
132 | queue.enqueue(x.key) |
||
133 | self.collect_keys(x.right, queue) |
||
134 | |||
135 | @staticmethod |
||
136 | def create(): |
||
137 | return BinarySearchTree() |
||
138 | |||
139 | @staticmethod |
||
140 | def create_red_black_tree(): |
||
141 | return RedBlackTree() |
||
142 | |||
143 | |||
144 | class RedBlackTree(BinarySearchTree): |
||
145 | def put(self, key, value): |
||
146 | self.root = self._put2(self.root, key, value) |
||
147 | |||
148 | def _put2(self, x, key, value): |
||
149 | if x is None: |
||
150 | return Node(key, value, 1) |
||
151 | compared = cmp(key, x.key) |
||
152 | if compared < 0: |
||
153 | x.left = self._put2(x.left, key, value) |
||
154 | elif compared > 0: |
||
155 | x.right = self._put2(x.right, key, value) |
||
156 | else: |
||
157 | x.value = value |
||
158 | |||
159 | if self.is_red(x.right) and not self.is_red(x.left): |
||
160 | x = self.rotate_left(x) |
||
161 | if self.is_red(x.left) and self.is_red(x.left.left): |
||
162 | x = self.rotate_right(x) |
||
163 | if self.is_red(x.left) and self.is_red(x.right): |
||
164 | x = self.flip_colors(x) |
||
165 | |||
166 | x.count = 1 + _count(x.left) + _count(x.right) |
||
167 | return x |
||
168 | |||
169 | def _delete(self, x, key): |
||
170 | if x is None: |
||
171 | return None |
||
172 | |||
173 | compared = cmp(key, x.key) |
||
174 | View Code Duplication | if compared < 0: |
|
0 ignored issues
–
show
|
|||
175 | x.left = self._delete(x.left, key) |
||
176 | elif compared > 0: |
||
177 | x.right = self._delete(x.right, key) |
||
178 | else: |
||
179 | if x.left is None: |
||
180 | return x.right |
||
181 | elif x.right is None: |
||
182 | return x.left |
||
183 | else: |
||
184 | m = self.min(x.right) |
||
185 | m.right = self.del_min(x.right) |
||
186 | m.left = x.left |
||
187 | |||
188 | x = m |
||
189 | |||
190 | if self.is_red(x.right) and not self.is_red(x.left): |
||
191 | x = self.rotate_left(x) |
||
192 | if self.is_red(x.left) and self.is_red(x.left.left): |
||
193 | x = self.rotate_right(x) |
||
194 | if self.is_red(x.left) and self.is_red(x.right): |
||
195 | x = self.flip_colors(x) |
||
196 | |||
197 | x.count = 1 + _count(x.left) + _count(x.right) |
||
198 | return x |
||
199 | |||
200 | def is_red(self, x): |
||
201 | if x is None: |
||
202 | return False |
||
203 | return x.red == 1 |
||
204 | |||
205 | def rotate_left(self, x): |
||
206 | m = x.right |
||
207 | x.right = m.left |
||
208 | m.left = x |
||
209 | |||
210 | m.red = x.red |
||
211 | x.red = 1 |
||
212 | |||
213 | x.count = 1 + _count(x.left) + _count(x.right) |
||
214 | return m |
||
215 | |||
216 | def rotate_right(self, x): |
||
217 | m = x.left |
||
218 | x.left = m.right |
||
219 | m.right = x |
||
220 | |||
221 | m.red = x.red |
||
222 | x.red = 1 |
||
223 | |||
224 | x.count = 1 + _count(x.left) + _count(x.right) |
||
225 | return m |
||
226 | |||
227 | def flip_colors(self, x): |
||
228 | if x.left is not None: |
||
229 | x.left.red = 0 |
||
230 | if x.right is not None: |
||
231 | x.right.red = 0 |
||
232 | x.red = 1 |
||
233 | return x |
||
234 | |||
235 |