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