signalis-eb/godot/addons/collections/bst.gd
2025-06-10 22:41:26 -04:00

138 lines
3.9 KiB
GDScript

class_name BST
class _Node extends RefCounted:
var value: int
var left: Option = Option.none
var right: Option = Option.none
func _init(value: int) -> void:
self.value = value
func size() -> int:
return BST.size(self)
func has_children() -> bool:
return BST.has_children(self)
func balance() -> _Node:
return BST.balance(self).unwrap()
func insert(value: int) -> _Node:
return BST.insert(self, value)
func search(value: int) -> Option:
return BST.search(self, value)
func remove(value: int) -> Option:
return BST.remove(self, value)
func iter() -> Iterator:
return BST.iter(self)
static func from(value: int) -> _Node:
return _Node.new(value)
static func size(node: _Node) -> int:
var left = node.left.map(size).or_default(0)
var right = node.right.map(size).or_default(0)
return left + right + 1
static func _store_in_order(_root: Option, nodes: Array[_Node] = []) -> Array[_Node]:
if _root.is_none(): return nodes
var root = _root.unwrap()
_store_in_order(root.left, nodes)
nodes.append(root)
_store_in_order(root.right, nodes)
return nodes
static func _build_balanced_tree(nodes: Array[_Node], start: int = 0, end: int = len(nodes) - 1) -> Option:
if start > end:
return Option.none
var mid: int = (start + end) / 2
var root = nodes[mid]
root.left = _build_balanced_tree(nodes, start, mid - 1)
root.right = _build_balanced_tree(nodes, mid + 1, end)
return Option.some(root)
static func balance(root: _Node) -> Option:
return _build_balanced_tree(_store_in_order(Option.some(root)))
static func has_children(node: _Node) -> bool:
return node.left.is_none() and node.right.is_none()
static func _insert(_root: Option, node: _Node) -> _Node:
if _root.is_none(): return node
var root = _root.unwrap()
if root.value == node.value: return root
elif root.value < node.value:
root.right = Option.some(_insert(root.right, node))
else:
root.left = Option.some(_insert(root.left, node))
return root
static func insert(root: _Node, value: int) -> _Node:
return _insert(Option.some(root), _Node.new(value))
static func _search(_root: Option, value: int) -> Option:
if _root.is_none():
return Option.none
var root = _root.unwrap()
if root.value == value:
return _root
if root.value < value:
return _search(root.right, value)
return _search(root.left, value)
static func search(root: _Node, value: int) -> Option:
return _search(Option.some(root), value)
static func _get_successor(node: _Node) -> _Node:
var current = node.right
while current.is_some() and current.left.is_some():
current = node.left.unwrap()
return current
static func _remove(_root: Option, value: int) -> Option:
if _root.is_none(): return _root
var root = _root.unwrap()
if root.value < value:
return _remove(root.right, value)
elif root.value > value:
return _remove(root.left, value)
else:
if root.right.is_none():
return root.left
elif root.left.is_none():
return root.right
var successor = _get_successor(root)
root.value = successor.value
root.right = remove(root.right, successor.value)
return Option.some(root)
static func remove(root: _Node, value: int) -> Option:
return _remove(Option.some(root), value)
static func traverse(root: _Node, fn: Callable):
root.left.map(fn.call)
fn.call(root.value)
root.right.map(fn.call)
static func iter(node: _Node) -> Iterator:
var iters = Option.collect_some([
node.left.map(iter),
Option.some(SingletonIterator.new(node.value)),
node.right.map(iter)
], true).unwrap()
var iter_arr: Array[Iterator]
iter_arr.assign(iters)
return FusedIterator.new(iter_arr)