blob: 81373fbfdf093747dd87c5edf67c7aa93ef17d59 [file] [log] [blame]
Matthias Andreas Benkard832a54e2019-01-29 09:27:38 +01001// Copyright 2010 Petar Maymounkov. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5// A Left-Leaning Red-Black (LLRB) implementation of 2-3 balanced binary search trees,
6// based on the following work:
7//
8// http://www.cs.princeton.edu/~rs/talks/LLRB/08Penn.pdf
9// http://www.cs.princeton.edu/~rs/talks/LLRB/LLRB.pdf
10// http://www.cs.princeton.edu/~rs/talks/LLRB/Java/RedBlackBST.java
11//
12// 2-3 trees (and the run-time equivalent 2-3-4 trees) are the de facto standard BST
13// algoritms found in implementations of Python, Java, and other libraries. The LLRB
14// implementation of 2-3 trees is a recent improvement on the traditional implementation,
15// observed and documented by Robert Sedgewick.
16//
17package llrb
18
19// Tree is a Left-Leaning Red-Black (LLRB) implementation of 2-3 trees
20type LLRB struct {
21 count int
22 root *Node
23}
24
25type Node struct {
26 Item
27 Left, Right *Node // Pointers to left and right child nodes
28 Black bool // If set, the color of the link (incoming from the parent) is black
29 // In the LLRB, new nodes are always red, hence the zero-value for node
30}
31
32type Item interface {
33 Less(than Item) bool
34}
35
36//
37func less(x, y Item) bool {
38 if x == pinf {
39 return false
40 }
41 if x == ninf {
42 return true
43 }
44 return x.Less(y)
45}
46
47// Inf returns an Item that is "bigger than" any other item, if sign is positive.
48// Otherwise it returns an Item that is "smaller than" any other item.
49func Inf(sign int) Item {
50 if sign == 0 {
51 panic("sign")
52 }
53 if sign > 0 {
54 return pinf
55 }
56 return ninf
57}
58
59var (
60 ninf = nInf{}
61 pinf = pInf{}
62)
63
64type nInf struct{}
65
66func (nInf) Less(Item) bool {
67 return true
68}
69
70type pInf struct{}
71
72func (pInf) Less(Item) bool {
73 return false
74}
75
76// New() allocates a new tree
77func New() *LLRB {
78 return &LLRB{}
79}
80
81// SetRoot sets the root node of the tree.
82// It is intended to be used by functions that deserialize the tree.
83func (t *LLRB) SetRoot(r *Node) {
84 t.root = r
85}
86
87// Root returns the root node of the tree.
88// It is intended to be used by functions that serialize the tree.
89func (t *LLRB) Root() *Node {
90 return t.root
91}
92
93// Len returns the number of nodes in the tree.
94func (t *LLRB) Len() int { return t.count }
95
96// Has returns true if the tree contains an element whose order is the same as that of key.
97func (t *LLRB) Has(key Item) bool {
98 return t.Get(key) != nil
99}
100
101// Get retrieves an element from the tree whose order is the same as that of key.
102func (t *LLRB) Get(key Item) Item {
103 h := t.root
104 for h != nil {
105 switch {
106 case less(key, h.Item):
107 h = h.Left
108 case less(h.Item, key):
109 h = h.Right
110 default:
111 return h.Item
112 }
113 }
114 return nil
115}
116
117// Min returns the minimum element in the tree.
118func (t *LLRB) Min() Item {
119 h := t.root
120 if h == nil {
121 return nil
122 }
123 for h.Left != nil {
124 h = h.Left
125 }
126 return h.Item
127}
128
129// Max returns the maximum element in the tree.
130func (t *LLRB) Max() Item {
131 h := t.root
132 if h == nil {
133 return nil
134 }
135 for h.Right != nil {
136 h = h.Right
137 }
138 return h.Item
139}
140
141func (t *LLRB) ReplaceOrInsertBulk(items ...Item) {
142 for _, i := range items {
143 t.ReplaceOrInsert(i)
144 }
145}
146
147func (t *LLRB) InsertNoReplaceBulk(items ...Item) {
148 for _, i := range items {
149 t.InsertNoReplace(i)
150 }
151}
152
153// ReplaceOrInsert inserts item into the tree. If an existing
154// element has the same order, it is removed from the tree and returned.
155func (t *LLRB) ReplaceOrInsert(item Item) Item {
156 if item == nil {
157 panic("inserting nil item")
158 }
159 var replaced Item
160 t.root, replaced = t.replaceOrInsert(t.root, item)
161 t.root.Black = true
162 if replaced == nil {
163 t.count++
164 }
165 return replaced
166}
167
168func (t *LLRB) replaceOrInsert(h *Node, item Item) (*Node, Item) {
169 if h == nil {
170 return newNode(item), nil
171 }
172
173 h = walkDownRot23(h)
174
175 var replaced Item
176 if less(item, h.Item) { // BUG
177 h.Left, replaced = t.replaceOrInsert(h.Left, item)
178 } else if less(h.Item, item) {
179 h.Right, replaced = t.replaceOrInsert(h.Right, item)
180 } else {
181 replaced, h.Item = h.Item, item
182 }
183
184 h = walkUpRot23(h)
185
186 return h, replaced
187}
188
189// InsertNoReplace inserts item into the tree. If an existing
190// element has the same order, both elements remain in the tree.
191func (t *LLRB) InsertNoReplace(item Item) {
192 if item == nil {
193 panic("inserting nil item")
194 }
195 t.root = t.insertNoReplace(t.root, item)
196 t.root.Black = true
197 t.count++
198}
199
200func (t *LLRB) insertNoReplace(h *Node, item Item) *Node {
201 if h == nil {
202 return newNode(item)
203 }
204
205 h = walkDownRot23(h)
206
207 if less(item, h.Item) {
208 h.Left = t.insertNoReplace(h.Left, item)
209 } else {
210 h.Right = t.insertNoReplace(h.Right, item)
211 }
212
213 return walkUpRot23(h)
214}
215
216// Rotation driver routines for 2-3 algorithm
217
218func walkDownRot23(h *Node) *Node { return h }
219
220func walkUpRot23(h *Node) *Node {
221 if isRed(h.Right) && !isRed(h.Left) {
222 h = rotateLeft(h)
223 }
224
225 if isRed(h.Left) && isRed(h.Left.Left) {
226 h = rotateRight(h)
227 }
228
229 if isRed(h.Left) && isRed(h.Right) {
230 flip(h)
231 }
232
233 return h
234}
235
236// Rotation driver routines for 2-3-4 algorithm
237
238func walkDownRot234(h *Node) *Node {
239 if isRed(h.Left) && isRed(h.Right) {
240 flip(h)
241 }
242
243 return h
244}
245
246func walkUpRot234(h *Node) *Node {
247 if isRed(h.Right) && !isRed(h.Left) {
248 h = rotateLeft(h)
249 }
250
251 if isRed(h.Left) && isRed(h.Left.Left) {
252 h = rotateRight(h)
253 }
254
255 return h
256}
257
258// DeleteMin deletes the minimum element in the tree and returns the
259// deleted item or nil otherwise.
260func (t *LLRB) DeleteMin() Item {
261 var deleted Item
262 t.root, deleted = deleteMin(t.root)
263 if t.root != nil {
264 t.root.Black = true
265 }
266 if deleted != nil {
267 t.count--
268 }
269 return deleted
270}
271
272// deleteMin code for LLRB 2-3 trees
273func deleteMin(h *Node) (*Node, Item) {
274 if h == nil {
275 return nil, nil
276 }
277 if h.Left == nil {
278 return nil, h.Item
279 }
280
281 if !isRed(h.Left) && !isRed(h.Left.Left) {
282 h = moveRedLeft(h)
283 }
284
285 var deleted Item
286 h.Left, deleted = deleteMin(h.Left)
287
288 return fixUp(h), deleted
289}
290
291// DeleteMax deletes the maximum element in the tree and returns
292// the deleted item or nil otherwise
293func (t *LLRB) DeleteMax() Item {
294 var deleted Item
295 t.root, deleted = deleteMax(t.root)
296 if t.root != nil {
297 t.root.Black = true
298 }
299 if deleted != nil {
300 t.count--
301 }
302 return deleted
303}
304
305func deleteMax(h *Node) (*Node, Item) {
306 if h == nil {
307 return nil, nil
308 }
309 if isRed(h.Left) {
310 h = rotateRight(h)
311 }
312 if h.Right == nil {
313 return nil, h.Item
314 }
315 if !isRed(h.Right) && !isRed(h.Right.Left) {
316 h = moveRedRight(h)
317 }
318 var deleted Item
319 h.Right, deleted = deleteMax(h.Right)
320
321 return fixUp(h), deleted
322}
323
324// Delete deletes an item from the tree whose key equals key.
325// The deleted item is return, otherwise nil is returned.
326func (t *LLRB) Delete(key Item) Item {
327 var deleted Item
328 t.root, deleted = t.delete(t.root, key)
329 if t.root != nil {
330 t.root.Black = true
331 }
332 if deleted != nil {
333 t.count--
334 }
335 return deleted
336}
337
338func (t *LLRB) delete(h *Node, item Item) (*Node, Item) {
339 var deleted Item
340 if h == nil {
341 return nil, nil
342 }
343 if less(item, h.Item) {
344 if h.Left == nil { // item not present. Nothing to delete
345 return h, nil
346 }
347 if !isRed(h.Left) && !isRed(h.Left.Left) {
348 h = moveRedLeft(h)
349 }
350 h.Left, deleted = t.delete(h.Left, item)
351 } else {
352 if isRed(h.Left) {
353 h = rotateRight(h)
354 }
355 // If @item equals @h.Item and no right children at @h
356 if !less(h.Item, item) && h.Right == nil {
357 return nil, h.Item
358 }
359 // PETAR: Added 'h.Right != nil' below
360 if h.Right != nil && !isRed(h.Right) && !isRed(h.Right.Left) {
361 h = moveRedRight(h)
362 }
363 // If @item equals @h.Item, and (from above) 'h.Right != nil'
364 if !less(h.Item, item) {
365 var subDeleted Item
366 h.Right, subDeleted = deleteMin(h.Right)
367 if subDeleted == nil {
368 panic("logic")
369 }
370 deleted, h.Item = h.Item, subDeleted
371 } else { // Else, @item is bigger than @h.Item
372 h.Right, deleted = t.delete(h.Right, item)
373 }
374 }
375
376 return fixUp(h), deleted
377}
378
379// Internal node manipulation routines
380
381func newNode(item Item) *Node { return &Node{Item: item} }
382
383func isRed(h *Node) bool {
384 if h == nil {
385 return false
386 }
387 return !h.Black
388}
389
390func rotateLeft(h *Node) *Node {
391 x := h.Right
392 if x.Black {
393 panic("rotating a black link")
394 }
395 h.Right = x.Left
396 x.Left = h
397 x.Black = h.Black
398 h.Black = false
399 return x
400}
401
402func rotateRight(h *Node) *Node {
403 x := h.Left
404 if x.Black {
405 panic("rotating a black link")
406 }
407 h.Left = x.Right
408 x.Right = h
409 x.Black = h.Black
410 h.Black = false
411 return x
412}
413
414// REQUIRE: Left and Right children must be present
415func flip(h *Node) {
416 h.Black = !h.Black
417 h.Left.Black = !h.Left.Black
418 h.Right.Black = !h.Right.Black
419}
420
421// REQUIRE: Left and Right children must be present
422func moveRedLeft(h *Node) *Node {
423 flip(h)
424 if isRed(h.Right.Left) {
425 h.Right = rotateRight(h.Right)
426 h = rotateLeft(h)
427 flip(h)
428 }
429 return h
430}
431
432// REQUIRE: Left and Right children must be present
433func moveRedRight(h *Node) *Node {
434 flip(h)
435 if isRed(h.Left.Left) {
436 h = rotateRight(h)
437 flip(h)
438 }
439 return h
440}
441
442func fixUp(h *Node) *Node {
443 if isRed(h.Right) {
444 h = rotateLeft(h)
445 }
446
447 if isRed(h.Left) && isRed(h.Left.Left) {
448 h = rotateRight(h)
449 }
450
451 if isRed(h.Left) && isRed(h.Right) {
452 flip(h)
453 }
454
455 return h
456}