feat: Add btree with insertion, deletion, and search (#703)

* feat: Add btree with insertion, deletion, and search

* Lazy allocate root and handle cases with null root

* Fix formatting

---------

Co-authored-by: Rak Laptudirm <[email protected]>

Author: caojoshua

structure/tree/btree.go

@@ -0,0 +1,355 @@
+// B-tree is a self balancing tree that promotes data locality.
+// For more details, see https://en.wikipedia.org/wiki/B-tree
+
+package tree
+
+import "github.com/TheAlgorithms/Go/constraints"
+
+type BTreeNode[T constraints.Ordered] struct {
+	keys     []T
+	children []*BTreeNode[T]
+	numKeys  int
+	isLeaf   bool
+}
+
+type BTree[T constraints.Ordered] struct {
+	root    *BTreeNode[T]
+	maxKeys int
+}
+
+func minKeys(maxKeys int) int {
+	return (maxKeys - 1) / 2
+}
+
+func NewBTreeNode[T constraints.Ordered](maxKeys int, isLeaf bool) *BTreeNode[T] {
+	if maxKeys <= 0 {
+		panic("BTree maxKeys cannot be zero")
+	}
+	return &BTreeNode[T]{
+		keys:     make([]T, maxKeys),
+		children: make([]*BTreeNode[T], maxKeys+1),
+		isLeaf:   isLeaf,
+	}
+}
+
+func NewBTree[T constraints.Ordered](maxKeys int) *BTree[T] {
+	if maxKeys <= 2 {
+		panic("Must be >= 3 keys")
+	}
+	return &BTree[T]{
+		root:    nil,
+		maxKeys: maxKeys,
+	}
+}
+
+func (node *BTreeNode[T]) Verify(tree *BTree[T]) {
+	minKeys := minKeys(tree.maxKeys)
+	if node != tree.root && node.numKeys < minKeys {
+		panic("node has too few keys")
+	} else if node.numKeys > tree.maxKeys {
+		panic("node has too many keys")
+	}
+}
+
+func (node *BTreeNode[T]) IsFull(maxKeys int) bool {
+	return node.numKeys == maxKeys
+}
+
+func (node *BTreeNode[T]) Search(key T) bool {
+	i := 0
+	for ; i < node.numKeys; i++ {
+		if key == node.keys[i] {
+			return true
+		}
+		if key < node.keys[i] {
+			break
+		}
+	}
+	if node.isLeaf {
+		return false
+	}
+	return node.children[i].Search(key)
+}
+
+func (tree *BTree[T]) Search(key T) bool {
+	if tree.root == nil {
+		return false
+	}
+	return tree.root.Search(key)
+}
+
+func (node *BTreeNode[T]) InsertKeyChild(key T, child *BTreeNode[T]) {
+	i := node.numKeys
+	node.children[i+1] = node.children[i]
+	for ; i > 0; i-- {
+		if key > node.keys[i-1] {
+			node.keys[i] = key
+			node.children[i] = child
+			break
+		}
+		node.keys[i] = node.keys[i-1]
+		node.children[i] = node.children[i-1]
+	}
+	if i == 0 {
+		node.keys[0] = key
+		node.children[0] = child
+	}
+	node.numKeys++
+}
+
+func (node *BTreeNode[T]) Append(key T, child *BTreeNode[T]) {
+	node.keys[node.numKeys] = key
+	node.children[node.numKeys+1] = child
+	node.numKeys++
+}
+
+// Add all of other's keys starting from idx and children starting from idx + 1
+func (node *BTreeNode[T]) Concat(other *BTreeNode[T], idx int) {
+	for i := 0; i < other.numKeys-idx; i++ {
+		node.keys[node.numKeys+i] = other.keys[i+idx]
+		node.children[node.numKeys+i+1] = other.children[i+idx+1]
+	}
+	node.numKeys += other.numKeys - idx
+}
+
+// Transform:
+//
+//	A B
+//	 |
+//
+// a b c d
+//
+// Into:
+//
+//	A c B
+//	 / \
+//
+// a b  d
+func (parent *BTreeNode[T]) Split(idx int, maxKeys int) {
+	child := parent.children[idx]
+	midKeyIndex := maxKeys / 2
+	rightChild := NewBTreeNode[T](maxKeys, child.isLeaf)
+	rightChild.Concat(child, midKeyIndex+1)
+	rightChild.children[0] = child.children[midKeyIndex+1]
+
+	// Reuse child as the left node
+	child.numKeys = midKeyIndex
+
+	// Insert the child's mid index to the parent
+	for i := parent.numKeys; i > idx; i-- {
+		parent.keys[i] = parent.keys[i-1]
+		parent.children[i+1] = parent.children[i]
+	}
+	parent.keys[idx] = child.keys[midKeyIndex]
+	parent.children[idx] = child
+	parent.children[idx+1] = rightChild
+	parent.numKeys += 1
+}
+
+func (node *BTreeNode[T]) InsertNonFull(tree *BTree[T], key T) {
+	node.Verify(tree)
+	if node.IsFull(tree.maxKeys) {
+		panic("Called InsertNonFull() with a full node")
+	}
+
+	if node.isLeaf {
+		// Node is a leaf. Directly insert the key.
+		node.InsertKeyChild(key, nil)
+		return
+	}
+
+	// Find the child node to insert into
+	i := 0
+	for ; i < node.numKeys; i++ {
+		if key < node.keys[i] {
+			break
+		}
+	}
+
+	if node.children[i].IsFull(tree.maxKeys) {
+		node.Split(i, tree.maxKeys)
+		if key > node.keys[i] {
+			i++
+		}
+	}
+	node.children[i].InsertNonFull(tree, key)
+}
+
+func (tree *BTree[T]) Insert(key T) {
+	if tree.root == nil {
+		tree.root = NewBTreeNode[T](tree.maxKeys, true)
+		tree.root.keys[0] = key
+		tree.root.numKeys = 1
+		return
+	}
+
+	if tree.root.IsFull(tree.maxKeys) {
+		newRoot := NewBTreeNode[T](tree.maxKeys, false)
+		newRoot.numKeys = 0
+		newRoot.children[0] = tree.root
+		newRoot.Split(0, tree.maxKeys)
+		tree.root = newRoot
+	}
+	tree.root.InsertNonFull(tree, key)
+}
+
+func (node *BTreeNode[T]) DeleteIthKey(i int) {
+	if i >= node.numKeys {
+		panic("deleting out of bounds key")
+	}
+	for j := i; j < node.numKeys-1; j++ {
+		node.keys[j] = node.keys[j+1]
+		node.children[j+1] = node.children[j+2]
+	}
+	node.numKeys--
+}
+
+// Transform:
+//
+//	A B C
+//	 / \
+//	a   b
+//
+// Into:
+//
+//	A C
+//	 |
+//
+// a B c
+func (node *BTreeNode[T]) Merge(idx int) {
+	if node.isLeaf {
+		panic("cannot merge when leaf node is parent")
+	}
+	left := node.children[idx]
+	right := node.children[idx+1]
+	left.Append(node.keys[idx], right.children[0])
+	left.Concat(right, 0)
+	node.DeleteIthKey(idx)
+}
+
+func (node *BTreeNode[T]) Min() T {
+	if node.isLeaf {
+		return node.keys[0]
+	}
+	return node.children[0].Min()
+}
+
+func (node *BTreeNode[T]) Max() T {
+	if node.isLeaf {
+		return node.keys[node.numKeys-1]
+	}
+	return node.children[node.numKeys].Max()
+}
+
+func (node *BTreeNode[T]) Delete(tree *BTree[T], key T) {
+	node.Verify(tree)
+	if node.isLeaf {
+		// Case 1: Node is a leaf. Directly delete the key.
+		for i := 0; i < node.numKeys; i++ {
+			if key == node.keys[i] {
+				node.DeleteIthKey(i)
+				return
+			}
+		}
+		return
+	}
+
+	minKeys := minKeys(tree.maxKeys)
+	i := 0
+	for ; i < node.numKeys; i++ {
+		if key == node.keys[i] {
+			// Case 2: key exists in a non-leaf node
+			left := node.children[i]
+			right := node.children[i+1]
+			if left.numKeys > minKeys {
+				// Replace the key we want to delete with the max key from the left
+				// subtree. Then delete that key from the left subtree.
+				//  A B C
+				//   /
+				// a b c
+				//
+				// If we want to delete `B`, then replace `B` with `c`, and delete `c` in the subtree.
+				//  A c C
+				//   /
+				// a b
+				replacementKey := left.Max()
+				node.keys[i] = replacementKey
+				left.Delete(tree, replacementKey)
+			} else if right.numKeys > minKeys {
+				// Replace the key we want to delete with the min key from the right
+				// subtree. Then delete that key in the right subtree. Mirrors the
+				// transformation above for replacing from the left subtree.
+				replacementKey := right.Min()
+				node.keys[i] = replacementKey
+				right.Delete(tree, replacementKey)
+			} else {
+				// Both left and right subtrees have the minimum number of keys. Merge
+				// the left tree, the deleted key, and the right tree together into the
+				// left tree. Then recursively delete the key in the left tree.
+				if left.numKeys != minKeys || right.numKeys != minKeys {
+					panic("nodes should not have less than the minimum number of keys")
+				}
+				node.Merge(i)
+				left.Delete(tree, key)
+			}
+			return
+		}
+
+		if key < node.keys[i] {
+			break
+		}
+	}
+
+	// Case 3: key may exist in a child node.
+	child := node.children[i]
+	if child.numKeys == minKeys {
+		// Before we recurse into the child node, make sure it has more than
+		// the minimum number of keys.
+		if i > 0 && node.children[i-1].numKeys > minKeys {
+			// Take a key from the left sibling
+			// Transform:
+			//  A  B  C
+			//   /   \
+			// a b    c
+			//
+			// Into:
+			//  A  b  C
+			//   /   \
+			//  a    B c
+			left := node.children[i-1]
+			child.InsertKeyChild(node.keys[i-1], left.children[left.numKeys])
+			node.keys[i-1] = left.keys[left.numKeys-1]
+			left.numKeys--
+		} else if i < node.numKeys && node.children[i+1].numKeys > minKeys {
+			// Take a key from the right sibling. Mirrors the transformation above for taking a key from the left sibling.
+			right := node.children[i+1]
+			child.Append(node.keys[i], right.children[0])
+			node.keys[i] = right.keys[0]
+			right.children[0] = right.children[1]
+			right.DeleteIthKey(0)
+		} else {
+			if i == 0 {
+				// Merge with right sibling
+				node.Merge(i)
+			} else {
+				// Merge with left sibling
+				node.Merge(i - 1)
+				child = node.children[i-1]
+			}
+		}
+	}
+	if child.numKeys == minKeys {
+		panic("cannot delete key from node with minimum number of keys")
+	}
+	child.Delete(tree, key)
+}
+
+func (tree *BTree[T]) Delete(key T) {
+	if tree.root == nil {
+		return
+	}
+	tree.root.Delete(tree, key)
+	if tree.root.numKeys == 0 {
+		tree.root = tree.root.children[0]
+	}
+}