Initial commit

tree.go

@@ -0,0 +1,215 @@
+package stl
+
+import (
+	"math/rand/v2"
+)
+
+// Node is a node on the binary search tree.
+type Node[K comparable, V any, C Sorter[K]] struct {
+	lson, rson *Node[K, V, C]
+	parent     *Node[K, V, C]
+	bal        uint32 // for treap balancing
+
+	key   K
+	Value V // The value on the node.
+}
+
+// Key is read-only access to the key of the node.
+func (n *Node[K, V, C]) Key() K {
+	return n.key
+}
+
+// Tree is a binary search tree. Always take
+// addresses and don't work with values.
+//
+// The zero value is a usable, empty tree.
+type Tree[K comparable, V any, C Sorter[K]] struct {
+	root *Node[K, V, C]
+}
+
+// TreeInc is an alias for a tree compared by the
+// default Less operator.
+type TreeInc[K ImplicitSortable, V any] = Tree[K, V, Less[K]]
+
+// TreeDec is an alias for a tree compared by the
+// Greater operator.
+type TreeDec[K ImplicitSortable, V any] = Tree[K, V, Greater[K]]
+
+// Clear clears the tree, erasing every element on it.
+func (t *Tree[K, V, C]) Clear() {
+	t.root = nil
+}
+
+// Insert tries to insert a new node on the tree, locating
+// the value if the key is already in the tree.
+//
+// Returns the new (or found) value on the tree, and true/false
+// if the key was inserted or found.
+func (t *Tree[K, V, C]) Insert(key K, value V) (value_on_tree *V, is_added bool) {
+	node, is_added := t.InsertNode(key, value)
+	return &node.Value, is_added
+}
+
+// InsertNode does the same as Insert, but returns the entire *Node.
+func (t *Tree[K, V, C]) InsertNode(key K, value V) (node *Node[K, V, C], is_added bool) {
+	t.root = t.realInsertNode(t.root, nil, key, value, &node, &is_added)
+	return
+}
+
+// realInsertNode recursivly searches and tries to insert the new node.
+// If the key is found, it's passed into *result.
+func (t *Tree[K, V, C]) realInsertNode(cur, parent *Node[K, V, C], key K, value V, result **Node[K, V, C], is_added *bool) (node *Node[K, V, C]) {
+	var c C
+	if cur == nil {
+		*result = &Node[K, V, C]{
+			parent: parent,
+			key:    key,
+			Value:  value,
+			bal:    rand.Uint32(),
+		}
+		*is_added = true
+		return *result
+	} else if c.Compare(key, cur.key) {
+		// key < now.key
+		cur.lson = t.realInsertNode(cur.lson, cur, key, value, result, is_added)
+		return cur
+	} else if c.Compare(cur.key, key) {
+		// key > now.key
+		cur.rson = t.realInsertNode(cur.rson, cur, key, value, result, is_added)
+		return cur
+	} else {
+		// key == now.key
+		*result = cur
+		*is_added = false
+		return cur
+	}
+
+}
+
+// Find finds the value related to the key on the tree,
+// or nil if it is not found.
+func (t *Tree[K, V, C]) Find(key K) (value_on_tree *V) {
+	node := t.FindNode(key)
+	if node == nil {
+		return nil
+	} else {
+		return &node.Value
+	}
+}
+
+// FindNode finds the node instead, or nil if it is not found.
+func (t *Tree[K, V, C]) FindNode(key K) (node *Node[K, V, C]) {
+	return t.realFindNode(t.root, key)
+}
+
+func (t *Tree[K, V, C]) realFindNode(cur *Node[K, V, C], key K) (node *Node[K, V, C]) {
+	var c C
+	if cur == nil {
+		return nil
+	} else if c.Compare(key, cur.key) {
+		// key < cur.key
+		return t.realFindNode(cur.lson, key)
+	} else if c.Compare(cur.key, key) {
+		// key > cur.key
+		return t.realFindNode(cur.rson, key)
+	} else {
+		// key == cur.key
+		return cur
+	}
+}
+
+// Deletes a node, does nothing if node is Nil.
+//
+// Use in tandem with FindNode to delete a key.
+func (t *Tree[K, V, C]) Delete(node *Node[K, V, C]) {
+	if node == nil {
+		return
+	}
+
+	for node.lson != nil && node.rson != nil {
+		if node.lson.bal < node.rson.bal {
+			node.lson.rotate(&t.root)
+		} else {
+			node.rson.rotate(&t.root)
+		}
+	}
+
+	if node == t.root {
+		// select new root
+		if node.lson != nil {
+			t.root = node.lson
+		} else {
+			t.root = node.rson
+		}
+	}
+
+	if node.lson != nil {
+		node.parent.connect(node.lson, node.tell())
+	} else {
+		node.parent.connect(node.rson, node.tell())
+	}
+
+	// delete everything for the GC
+	// Seriously, why use a tree when you have a GC?
+	node.lson = nil
+	node.rson = nil
+	node.parent = nil
+	var k K
+	var v V
+	node.key = k
+	node.Value = v
+}
+
+func (t *Tree[K, V, C]) FirstNode() (node *Node[K, V, C]) {
+	node = t.root
+	if node == nil {
+		return nil
+	}
+	for node.lson != nil {
+		node = node.lson
+	}
+	return node
+}
+
+func (t *Tree[K, V, C]) LastNode() (node *Node[K, V, C]) {
+	node = t.root
+	if node == nil {
+		return nil
+	}
+	for node.rson != nil {
+		node = node.rson
+	}
+	return node
+}
+
+func (node *Node[K, V, C]) Next() (next *Node[K, V, C]) {
+	if node.rson != nil {
+		next = node.rson
+		for next.lson != nil {
+			next = next.lson
+		}
+		return
+	} else {
+		next = node
+		for next.parent != nil && next.tell() == treeRight {
+			next = next.parent
+		}
+		return next.parent
+	}
+}
+
+func (node *Node[K, V, C]) Previous() (prev *Node[K, V, C]) {
+	if node.lson != nil {
+		prev = node.lson
+		for prev.rson != nil {
+			prev = prev.rson
+		}
+		return prev
+	} else {
+		prev = node
+		for prev.parent != nil && prev.tell() == treeLeft {
+			prev = prev.parent
+		}
+		return prev.parent
+	}
+}