1package avl
2
3//----------------------------------------
4// Node
5
6// Node represents a node in an AVL tree.
7type Node struct {
8 key string // key is the unique identifier for the node.
9 value interface{} // value is the data stored in the node.
10 height int8 // height is the height of the node in the tree.
11 size int // size is the number of nodes in the subtree rooted at this node.
12 leftNode *Node // leftNode is the left child of the node.
13 rightNode *Node // rightNode is the right child of the node.
14}
15
16// NewNode creates a new node with the given key and value.
17func NewNode(key string, value interface{}) *Node {
18 return &Node{
19 key: key,
20 value: value,
21 height: 0,
22 size: 1,
23 }
24}
25
26// Size returns the size of the subtree rooted at the node.
27func (node *Node) Size() int {
28 if node == nil {
29 return 0
30 }
31 return node.size
32}
33
34// IsLeaf checks if the node is a leaf node (has no children).
35func (node *Node) IsLeaf() bool {
36 return node.height == 0
37}
38
39// Key returns the key of the node.
40func (node *Node) Key() string {
41 return node.key
42}
43
44// Value returns the value of the node.
45func (node *Node) Value() interface{} {
46 return node.value
47}
48
49// _copy creates a copy of the node (excluding value).
50func (node *Node) _copy() *Node {
51 if node.height == 0 {
52 panic("Why are you copying a value node?")
53 }
54 return &Node{
55 key: node.key,
56 height: node.height,
57 size: node.size,
58 leftNode: node.leftNode,
59 rightNode: node.rightNode,
60 }
61}
62
63// Has checks if a node with the given key exists in the subtree rooted at the node.
64func (node *Node) Has(key string) (has bool) {
65 if node == nil {
66 return false
67 }
68 if node.key == key {
69 return true
70 }
71 if node.height == 0 {
72 return false
73 }
74 if key < node.key {
75 return node.getLeftNode().Has(key)
76 }
77 return node.getRightNode().Has(key)
78}
79
80// Get searches for a node with the given key in the subtree rooted at the node
81// and returns its index, value, and whether it exists.
82func (node *Node) Get(key string) (index int, value interface{}, exists bool) {
83 if node == nil {
84 return 0, nil, false
85 }
86
87 if node.height == 0 {
88 if node.key == key {
89 return 0, node.value, true
90 }
91 if node.key < key {
92 return 1, nil, false
93 }
94 return 0, nil, false
95 }
96
97 if key < node.key {
98 return node.getLeftNode().Get(key)
99 }
100
101 rightNode := node.getRightNode()
102 index, value, exists = rightNode.Get(key)
103 index += node.size - rightNode.size
104 return index, value, exists
105}
106
107// GetByIndex retrieves the key-value pair of the node at the given index
108// in the subtree rooted at the node.
109func (node *Node) GetByIndex(index int) (key string, value interface{}) {
110 if node.height == 0 {
111 if index == 0 {
112 return node.key, node.value
113 }
114 panic("GetByIndex asked for invalid index")
115 }
116 // TODO: could improve this by storing the sizes
117 leftNode := node.getLeftNode()
118 if index < leftNode.size {
119 return leftNode.GetByIndex(index)
120 }
121 return node.getRightNode().GetByIndex(index - leftNode.size)
122}
123
124// Set inserts a new node with the given key-value pair into the subtree rooted at the node,
125// and returns the new root of the subtree and whether an existing node was updated.
126//
127// XXX consider a better way to do this... perhaps split Node from Node.
128func (node *Node) Set(key string, value interface{}) (newSelf *Node, updated bool) {
129 if node == nil {
130 return NewNode(key, value), false
131 }
132
133 if node.height == 0 {
134 return node.setLeaf(key, value)
135 }
136
137 node = node._copy()
138 if key < node.key {
139 node.leftNode, updated = node.getLeftNode().Set(key, value)
140 } else {
141 node.rightNode, updated = node.getRightNode().Set(key, value)
142 }
143
144 if updated {
145 return node, updated
146 }
147
148 node.calcHeightAndSize()
149 return node.balance(), updated
150}
151
152// setLeaf inserts a new leaf node with the given key-value pair into the subtree rooted at the node,
153// and returns the new root of the subtree and whether an existing node was updated.
154func (node *Node) setLeaf(key string, value interface{}) (newSelf *Node, updated bool) {
155 if key == node.key {
156 return NewNode(key, value), true
157 }
158
159 if key < node.key {
160 return &Node{
161 key: node.key,
162 height: 1,
163 size: 2,
164 leftNode: NewNode(key, value),
165 rightNode: node,
166 }, false
167 }
168
169 return &Node{
170 key: key,
171 height: 1,
172 size: 2,
173 leftNode: node,
174 rightNode: NewNode(key, value),
175 }, false
176}
177
178// Remove deletes the node with the given key from the subtree rooted at the node.
179// returns the new root of the subtree, the new leftmost leaf key (if changed),
180// the removed value and the removal was successful.
181func (node *Node) Remove(key string) (
182 newNode *Node, newKey string, value interface{}, removed bool,
183) {
184 if node == nil {
185 return nil, "", nil, false
186 }
187 if node.height == 0 {
188 if key == node.key {
189 return nil, "", node.value, true
190 }
191 return node, "", nil, false
192 }
193 if key < node.key {
194 var newLeftNode *Node
195 newLeftNode, newKey, value, removed = node.getLeftNode().Remove(key)
196 if !removed {
197 return node, "", value, false
198 }
199 if newLeftNode == nil { // left node held value, was removed
200 return node.rightNode, node.key, value, true
201 }
202 node = node._copy()
203 node.leftNode = newLeftNode
204 node.calcHeightAndSize()
205 node = node.balance()
206 return node, newKey, value, true
207 }
208
209 var newRightNode *Node
210 newRightNode, newKey, value, removed = node.getRightNode().Remove(key)
211 if !removed {
212 return node, "", value, false
213 }
214 if newRightNode == nil { // right node held value, was removed
215 return node.leftNode, "", value, true
216 }
217 node = node._copy()
218 node.rightNode = newRightNode
219 if newKey != "" {
220 node.key = newKey
221 }
222 node.calcHeightAndSize()
223 node = node.balance()
224 return node, "", value, true
225}
226
227// getLeftNode returns the left child of the node.
228func (node *Node) getLeftNode() *Node {
229 return node.leftNode
230}
231
232// getRightNode returns the right child of the node.
233func (node *Node) getRightNode() *Node {
234 return node.rightNode
235}
236
237// rotateRight performs a right rotation on the node and returns the new root.
238// NOTE: overwrites node
239// TODO: optimize balance & rotate
240func (node *Node) rotateRight() *Node {
241 node = node._copy()
242 l := node.getLeftNode()
243 _l := l._copy()
244
245 _lrCached := _l.rightNode
246 _l.rightNode = node
247 node.leftNode = _lrCached
248
249 node.calcHeightAndSize()
250 _l.calcHeightAndSize()
251
252 return _l
253}
254
255// rotateLeft performs a left rotation on the node and returns the new root.
256// NOTE: overwrites node
257// TODO: optimize balance & rotate
258func (node *Node) rotateLeft() *Node {
259 node = node._copy()
260 r := node.getRightNode()
261 _r := r._copy()
262
263 _rlCached := _r.leftNode
264 _r.leftNode = node
265 node.rightNode = _rlCached
266
267 node.calcHeightAndSize()
268 _r.calcHeightAndSize()
269
270 return _r
271}
272
273// calcHeightAndSize updates the height and size of the node based on its children.
274// NOTE: mutates height and size
275func (node *Node) calcHeightAndSize() {
276 node.height = maxInt8(node.getLeftNode().height, node.getRightNode().height) + 1
277 node.size = node.getLeftNode().size + node.getRightNode().size
278}
279
280// calcBalance calculates the balance factor of the node.
281func (node *Node) calcBalance() int {
282 return int(node.getLeftNode().height) - int(node.getRightNode().height)
283}
284
285// balance balances the subtree rooted at the node and returns the new root.
286// NOTE: assumes that node can be modified
287// TODO: optimize balance & rotate
288func (node *Node) balance() (newSelf *Node) {
289 balance := node.calcBalance()
290 if balance >= -1 {
291 return node
292 }
293 if balance > 1 {
294 if node.getLeftNode().calcBalance() >= 0 {
295 // Left Left Case
296 return node.rotateRight()
297 }
298 // Left Right Case
299 left := node.getLeftNode()
300 node.leftNode = left.rotateLeft()
301 return node.rotateRight()
302 }
303
304 if node.getRightNode().calcBalance() <= 0 {
305 // Right Right Case
306 return node.rotateLeft()
307 }
308
309 // Right Left Case
310 right := node.getRightNode()
311 node.rightNode = right.rotateRight()
312 return node.rotateLeft()
313}
314
315// Shortcut for TraverseInRange.
316func (node *Node) Iterate(start, end string, cb func(*Node) bool) bool {
317 return node.TraverseInRange(start, end, true, true, cb)
318}
319
320// Shortcut for TraverseInRange.
321func (node *Node) ReverseIterate(start, end string, cb func(*Node) bool) bool {
322 return node.TraverseInRange(start, end, false, true, cb)
323}
324
325// TraverseInRange traverses all nodes, including inner nodes.
326// Start is inclusive and end is exclusive when ascending,
327// Start and end are inclusive when descending.
328// Empty start and empty end denote no start and no end.
329// If leavesOnly is true, only visit leaf nodes.
330// NOTE: To simulate an exclusive reverse traversal,
331// just append 0x00 to start.
332func (node *Node) TraverseInRange(start, end string, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
333 if node == nil {
334 return false
335 }
336 afterStart := (start == "" || start < node.key)
337 startOrAfter := (start == "" || start <= node.key)
338 beforeEnd := false
339 if ascending {
340 beforeEnd = (end == "" || node.key < end)
341 } else {
342 beforeEnd = (end == "" || node.key <= end)
343 }
344
345 // Run callback per inner/leaf node.
346 stop := false
347 if (!node.IsLeaf() && !leavesOnly) ||
348 (node.IsLeaf() && startOrAfter && beforeEnd) {
349 stop = cb(node)
350 if stop {
351 return stop
352 }
353 }
354 if node.IsLeaf() {
355 return stop
356 }
357
358 if ascending {
359 // check lower nodes, then higher
360 if afterStart {
361 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
362 }
363 if stop {
364 return stop
365 }
366 if beforeEnd {
367 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
368 }
369 } else {
370 // check the higher nodes first
371 if beforeEnd {
372 stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
373 }
374 if stop {
375 return stop
376 }
377 if afterStart {
378 stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)
379 }
380 }
381
382 return stop
383}
384
385// TraverseByOffset traverses all nodes, including inner nodes.
386// A limit of math.MaxInt means no limit.
387func (node *Node) TraverseByOffset(offset, limit int, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
388 if node == nil {
389 return false
390 }
391
392 // fast paths. these happen only if TraverseByOffset is called directly on a leaf.
393 if limit <= 0 || offset >= node.size {
394 return false
395 }
396 if node.IsLeaf() {
397 if offset > 0 {
398 return false
399 }
400 return cb(node)
401 }
402
403 // go to the actual recursive function.
404 return node.traverseByOffset(offset, limit, ascending, leavesOnly, cb)
405}
406
407// TraverseByOffset traverses the subtree rooted at the node by offset and limit,
408// in either ascending or descending order, and applies the callback function to each traversed node.
409// If leavesOnly is true, only leaf nodes are visited.
410func (node *Node) traverseByOffset(offset, limit int, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {
411 // caller guarantees: offset < node.size; limit > 0.
412 if !leavesOnly {
413 if cb(node) {
414 return true // Stop traversal if callback returns true
415 }
416 }
417 first, second := node.getLeftNode(), node.getRightNode()
418 if !ascending {
419 first, second = second, first
420 }
421 if first.IsLeaf() {
422 // either run or skip, based on offset
423 if offset > 0 {
424 offset--
425 } else {
426 if cb(first) {
427 return true // Stop traversal if callback returns true
428 }
429 limit--
430 if limit <= 0 {
431 return true // Stop traversal when limit is reached
432 }
433 }
434 } else {
435 // possible cases:
436 // 1 the offset given skips the first node entirely
437 // 2 the offset skips none or part of the first node, but the limit requires some of the second node.
438 // 3 the offset skips none or part of the first node, and the limit stops our search on the first node.
439 if offset >= first.size {
440 offset -= first.size // 1
441 } else {
442 if first.traverseByOffset(offset, limit, ascending, leavesOnly, cb) {
443 return true
444 }
445 // number of leaves which could actually be called from inside
446 delta := first.size - offset
447 offset = 0
448 if delta >= limit {
449 return true // 3
450 }
451 limit -= delta // 2
452 }
453 }
454
455 // because of the caller guarantees and the way we handle the first node,
456 // at this point we know that limit > 0 and there must be some values in
457 // this second node that we include.
458
459 // => if the second node is a leaf, it has to be included.
460 if second.IsLeaf() {
461 return cb(second)
462 }
463 // => if it is not a leaf, it will still be enough to recursively call this
464 // function with the updated offset and limit
465 return second.traverseByOffset(offset, limit, ascending, leavesOnly, cb)
466}
467
468// Only used in testing...
469func (node *Node) lmd() *Node {
470 if node.height == 0 {
471 return node
472 }
473 return node.getLeftNode().lmd()
474}
475
476// Only used in testing...
477func (node *Node) rmd() *Node {
478 if node.height == 0 {
479 return node
480 }
481 return node.getRightNode().rmd()
482}
483
484func maxInt8(a, b int8) int8 {
485 if a > b {
486 return a
487 }
488 return b
489}
node.gno
12.54 Kb ยท 489 lines
