gno.land/p/samcrew/avl⧉

1package avl23//----------------------------------------4// Node56// 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}1516// 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}2526// Size returns the size of the subtree rooted at the node.27func (node *Node) Size() int {28	if node == nil {29		return 030	}31	return node.size32}3334// IsLeaf checks if the node is a leaf node (has no children).35func (node *Node) IsLeaf() bool {36	return node.height == 037}3839// Key returns the key of the node.40func (node *Node) Key() string {41	return node.key42}4344// Value returns the value of the node.45func (node *Node) Value() any {46	return node.value47}4849func (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}6162// 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 false66	}67	if node.key == key {68		return true69	}70	if node.height == 0 {71		return false72	} else {73		if key < node.key {74			return node.getLeftNode().Has(key)75		} else {76			return node.getRightNode().Has(key)77		}78	}79}8081// Get searches for a node with the given key in the subtree rooted at the node82// 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, false86	}8788	if node.height == 0 {89		if node.key == key {90			return 0, node.value, true91		} else if node.key < key {92			return 1, nil, false93		} else {94			return 0, nil, false95		}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.size103			return index, value, exists104		}105	}106}107108// GetByIndex retrieves the key-value pair of the node at the given index109// in the subtree rooted at the node.110func (node *Node) GetByIndex(index int) (key string, value any) {111	if node.height == 0 {112		if index == 0 {113			return node.key, node.value114		} else {115			panic("GetByIndex asked for invalid index")116		}117	} else {118		// TODO: could improve this by storing the sizes119		leftNode := node.getLeftNode()120		if index < leftNode.size {121			return leftNode.GetByIndex(index)122		} else {123			return node.getRightNode().GetByIndex(index - leftNode.size)124		}125	}126}127128// Set inserts a new node with the given key-value pair into the subtree rooted at the node,129// and returns the new root of the subtree and whether an existing node was updated.130//131// XXX consider a better way to do this... perhaps split Node from Node.132func (node *Node) Set(key string, value any) (newSelf *Node, updated bool) {133	if node == nil {134		return NewNode(key, value), false135	}136	if node.height == 0 {137		if key < node.key {138			return &Node{139				key:       node.key,140				height:    1,141				size:      2,142				leftNode:  NewNode(key, value),143				rightNode: node,144			}, false145		} else if key == node.key {146			return NewNode(key, value), true147		} else {148			return &Node{149				key:       key,150				height:    1,151				size:      2,152				leftNode:  node,153				rightNode: NewNode(key, value),154			}, false155		}156	} else {157		node = node._copy()158		if key < node.key {159			node.leftNode, updated = node.getLeftNode().Set(key, value)160		} else {161			node.rightNode, updated = node.getRightNode().Set(key, value)162		}163		if updated {164			return node, updated165		} else {166			node.calcHeightAndSize()167			return node.balance(), updated168		}169	}170}171172// Remove deletes the node with the given key from the subtree rooted at the node.173// returns the new root of the subtree, the new leftmost leaf key (if changed),174// the removed value and the removal was successful.175func (node *Node) Remove(key string) (176	newNode *Node, newKey string, value any, removed bool,177) {178	if node == nil {179		return nil, "", nil, false180	}181	if node.height == 0 {182		if key == node.key {183			return nil, "", node.value, true184		} else {185			return node, "", nil, false186		}187	} else {188		if key < node.key {189			var newLeftNode *Node190			newLeftNode, newKey, value, removed = node.getLeftNode().Remove(key)191			if !removed {192				return node, "", value, false193			} else if newLeftNode == nil { // left node held value, was removed194				return node.rightNode, node.key, value, true195			}196			node = node._copy()197			node.leftNode = newLeftNode198			node.calcHeightAndSize()199			node = node.balance()200			return node, newKey, value, true201		} else {202			var newRightNode *Node203			newRightNode, newKey, value, removed = node.getRightNode().Remove(key)204			if !removed {205				return node, "", value, false206			} else if newRightNode == nil { // right node held value, was removed207				return node.leftNode, "", value, true208			}209			node = node._copy()210			node.rightNode = newRightNode211			if newKey != "" {212				node.key = newKey213			}214			node.calcHeightAndSize()215			node = node.balance()216			return node, "", value, true217		}218	}219}220221func (node *Node) getLeftNode() *Node {222	return node.leftNode223}224225func (node *Node) getRightNode() *Node {226	return node.rightNode227}228229// rotateRight performs a right rotation on the node and returns the new root.230// NOTE: overwrites node231// TODO: optimize balance & rotate232func (node *Node) rotateRight() *Node {233	node = node._copy()234	l := node.getLeftNode()235	_l := l._copy()236237	_lrCached := _l.rightNode238	_l.rightNode = node239	node.leftNode = _lrCached240241	node.calcHeightAndSize()242	_l.calcHeightAndSize()243244	return _l245}246247// rotateLeft performs a left rotation on the node and returns the new root.248// NOTE: overwrites node249// TODO: optimize balance & rotate250func (node *Node) rotateLeft() *Node {251	node = node._copy()252	r := node.getRightNode()253	_r := r._copy()254255	_rlCached := _r.leftNode256	_r.leftNode = node257	node.rightNode = _rlCached258259	node.calcHeightAndSize()260	_r.calcHeightAndSize()261262	return _r263}264265// calcHeightAndSize updates the height and size of the node based on its children.266// NOTE: mutates height and size267func (node *Node) calcHeightAndSize() {268	node.height = maxInt8(node.getLeftNode().height, node.getRightNode().height) + 1269	node.size = node.getLeftNode().size + node.getRightNode().size270}271272// calcBalance calculates the balance factor of the node.273func (node *Node) calcBalance() int {274	return int(node.getLeftNode().height) - int(node.getRightNode().height)275}276277// balance balances the subtree rooted at the node and returns the new root.278// NOTE: assumes that node can be modified279// TODO: optimize balance & rotate280func (node *Node) balance() (newSelf *Node) {281	balance := node.calcBalance()282	if balance > 1 {283		if node.getLeftNode().calcBalance() >= 0 {284			// Left Left Case285			return node.rotateRight()286		} else {287			// Left Right Case288			left := node.getLeftNode()289			node.leftNode = left.rotateLeft()290			return node.rotateRight()291		}292	}293	if balance < -1 {294		if node.getRightNode().calcBalance() <= 0 {295			// Right Right Case296			return node.rotateLeft()297		} else {298			// Right Left Case299			right := node.getRightNode()300			node.rightNode = right.rotateRight()301			return node.rotateLeft()302		}303	}304	// Nothing changed305	return node306}307308// Shortcut for TraverseInRange.309func (node *Node) Iterate(start, end string, cb func(*Node) bool) bool {310	return node.TraverseInRange(start, end, true, true, cb)311}312313// Shortcut for TraverseInRange.314func (node *Node) ReverseIterate(start, end string, cb func(*Node) bool) bool {315	return node.TraverseInRange(start, end, false, true, cb)316}317318// TraverseInRange traverses all nodes, including inner nodes.319// Start is inclusive and end is exclusive when ascending,320// Start and end are inclusive when descending.321// Empty start and empty end denote no start and no end.322// If leavesOnly is true, only visit leaf nodes.323// NOTE: To simulate an exclusive reverse traversal,324// just append 0x00 to start.325func (node *Node) TraverseInRange(start, end string, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {326	if node == nil {327		return false328	}329	afterStart := (start == "" || start < node.key)330	startOrAfter := (start == "" || start <= node.key)331	beforeEnd := false332	if ascending {333		beforeEnd = (end == "" || node.key < end)334	} else {335		beforeEnd = (end == "" || node.key <= end)336	}337338	// Run callback per inner/leaf node.339	stop := false340	if (!node.IsLeaf() && !leavesOnly) ||341		(node.IsLeaf() && startOrAfter && beforeEnd) {342		stop = cb(node)343		if stop {344			return stop345		}346	}347	if node.IsLeaf() {348		return stop349	}350351	if ascending {352		// check lower nodes, then higher353		if afterStart {354			stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)355		}356		if stop {357			return stop358		}359		if beforeEnd {360			stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)361		}362	} else {363		// check the higher nodes first364		if beforeEnd {365			stop = node.getRightNode().TraverseInRange(start, end, ascending, leavesOnly, cb)366		}367		if stop {368			return stop369		}370		if afterStart {371			stop = node.getLeftNode().TraverseInRange(start, end, ascending, leavesOnly, cb)372		}373	}374375	return stop376}377378// TraverseByOffset traverses all nodes, including inner nodes.379// A limit of math.MaxInt means no limit.380func (node *Node) TraverseByOffset(offset, limit int, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {381	if node == nil {382		return false383	}384385	// fast paths. these happen only if TraverseByOffset is called directly on a leaf.386	if limit <= 0 || offset >= node.size {387		return false388	}389	if node.IsLeaf() {390		if offset > 0 {391			return false392		}393		return cb(node)394	}395396	// go to the actual recursive function.397	return node.traverseByOffset(offset, limit, ascending, leavesOnly, cb)398}399400// TraverseByOffset traverses the subtree rooted at the node by offset and limit,401// in either ascending or descending order, and applies the callback function to each traversed node.402// If leavesOnly is true, only leaf nodes are visited.403func (node *Node) traverseByOffset(offset, limit int, ascending bool, leavesOnly bool, cb func(*Node) bool) bool {404	// caller guarantees: offset < node.size; limit > 0.405	if !leavesOnly {406		if cb(node) {407			return true // Stop traversal if callback returns true408		}409	}410	first, second := node.getLeftNode(), node.getRightNode()411	if !ascending {412		first, second = second, first413	}414	if first.IsLeaf() {415		// either run or skip, based on offset416		if offset > 0 {417			offset--418		} else {419			if cb(first) {420				return true // Stop traversal if callback returns true421			}422			limit--423			if limit <= 0 {424				return true // Stop traversal when limit is reached425			}426		}427	} else {428		// possible cases:429		// 1 the offset given skips the first node entirely430		// 2 the offset skips none or part of the first node, but the limit requires some of the second node.431		// 3 the offset skips none or part of the first node, and the limit stops our search on the first node.432		if offset >= first.size {433			offset -= first.size // 1434		} else {435			if first.traverseByOffset(offset, limit, ascending, leavesOnly, cb) {436				return true437			}438			// number of leaves which could actually be called from inside439			delta := first.size - offset440			offset = 0441			if delta >= limit {442				return true // 3443			}444			limit -= delta // 2445		}446	}447448	// because of the caller guarantees and the way we handle the first node,449	// at this point we know that limit > 0 and there must be some values in450	// this second node that we include.451452	// => if the second node is a leaf, it has to be included.453	if second.IsLeaf() {454		return cb(second)455	}456	// => if it is not a leaf, it will still be enough to recursively call this457	// function with the updated offset and limit458	return second.traverseByOffset(offset, limit, ascending, leavesOnly, cb)459}460461// Only used in testing...462func (node *Node) lmd() *Node {463	if node.height == 0 {464		return node465	}466	return node.getLeftNode().lmd()467}468469// Only used in testing...470func (node *Node) rmd() *Node {471	if node.height == 0 {472		return node473	}474	return node.getRightNode().rmd()475}476477func maxInt8(a, b int8) int8 {478	if a > b {479		return a480	}481	return b482}483