Tree/Balanced.bruijn
# MIT License, Copyright (c) 2023 Marvin Borner
# Balanced AVL tree for numerical leafs, inspired by Rudy Matela's implementation
# TODO: Extend for arbitrary orderable types?
# TODO: Analyze performance bottlenecks
# TODO: More tests
:import std/Combinator .
:import std/List L
:import std/Logic .
:import std/Number .
:import std/Option .
:import std/Pair .
# error return that can't happen (in theory)
error Ω
# unwraps tree from option (only use if not empty!)
unwrap unwrap-or error ⧗ (Option BalancedTree) → BalancedTree
!‣ unwrap
# empty tree
empty none ⧗ (Option BalancedTree)
# returns height of tree
height map-or (-1) ^‣ ⧗ (Option BalancedTree) → Number
:test (height empty) ((-1))
:test (height (some ((+5) : ((+42) : (none : none))))) ((+5))
# constructs a tree with a label and no branches
node [[[(max (height 0) ++(height 2)) : (2 : (1 : 0))]]] ⧗ (Option BalancedTree) → Number → (Option BalancedTree) → BalancedTree
# constructs a leaf node
leaf [node none 0 none] ⧗ Number → BalancedTree
:test (leaf (+42)) (++(-1) : (none : ((+42) : none)))
# returns the label of a tree
label [^(~(~0))] ⧗ BalancedTree → Number
?‣ label
:test (?(leaf (+42))) ((+42))
# returns the left branch of a tree
left [^(~0)] ⧗ BalancedTree → (Option BalancedTree)
//‣ left
# returns the right branch of a tree
right [~(~(~0))] ⧗ BalancedTree → (Option BalancedTree)
\\‣ right
# returns the balancing factor of a tree
factor map-or (+0) d ⧗ (Option BalancedTree) → Number
d [(height //0) - (height \\0)]
:test (factor (some (leaf (+42)))) (++(-1))
rotate-ll [node //(!(//0)) ?(//0) (some (node \\(!(//0)))) ?0 \\0] ⧗ BalancedTree → BalancedTree
rotate-rr [node (some (node //0 ?0 //(!(\\0)))) ?(\\0) \\(!(\\0))] ⧗ BalancedTree → BalancedTree
rotate-lr [rotate-ll (node (some (rotate-rr !(//0))) ?0 \\0)] ⧗ BalancedTree → BalancedTree
rotate-rl [rotate-rr (node //0 ?0 (some (rotate-ll !(\\0))))] ⧗ BalancedTree → BalancedTree
# balances a tree
balance [go (factor 0)] ⧗ (Option BalancedTree) → (Option BalancedTree)
go [=?0 else ((0 >? (+1)) left ((0 <? (-1)) right else))]
left some (((factor //(!1)) =? (-1)) (rotate-lr !1) (rotate-ll !1))
right some (((factor \\(!1)) =? (+1)) (rotate-rl !1) (rotate-rr !1))
else 1
# inserts a number into a tree
insert z [[[rec]]] ⧗ Number → (Option BalancedTree) → (Option BalancedTree)
rec none? 0 (some (leaf 1)) (balance (u 1 ?(!0)))
u compare-case eq lt gt
eq 0
lt some (node (2 1 //(!0)) ?(!0) \\(!0))
gt some (node //(!0) ?(!0) (2 1 \\(!0)))
# returns true if a number is in a tree
has? z [[[rec]]] ⧗ Number → (Option BalancedTree) → Boolean
rec none? 0 false (u 1 ?(!0))
u compare-case eq lt gt
eq true
lt 2 1 //(!0)
gt 2 1 \\(!0)
:test (has? (+42) empty) (false)
# number of elements in tree (slow)
size z [[rec]] ⧗ (Option BalancedTree) → Number
rec none? 0 case-empty case-full
case-full (1 //(!0)) + (1 \\(!0)) + (+1)
case-empty (+0)
# converts a tree to a list
list! z [[map-or L.empty go 0]] ⧗ (Option BalancedTree) → (List Number)
go [L.append (L.append (2 //0) L.{}(?0)) (2 \\0)]
# converts a list to a tree
from-list L.foldr insert empty ⧗ (List Number) → (Option BalancedTree)