datatypes: add a binary search tree implementation (#13453)
parent
4a765bc33b
commit
7bd8503170
|
@ -0,0 +1,28 @@
|
||||||
|
import datatypes
|
||||||
|
|
||||||
|
struct KeyVal {
|
||||||
|
mut:
|
||||||
|
key int
|
||||||
|
val int
|
||||||
|
}
|
||||||
|
|
||||||
|
fn (a KeyVal) == (b KeyVal) bool {
|
||||||
|
return a.key == b.key
|
||||||
|
}
|
||||||
|
|
||||||
|
fn (a KeyVal) < (b KeyVal) bool {
|
||||||
|
return a.key < b.key
|
||||||
|
}
|
||||||
|
|
||||||
|
fn main() {
|
||||||
|
mut bst := datatypes.BSTree<KeyVal>{}
|
||||||
|
bst.insert(KeyVal{ key: 1, val: 12 })
|
||||||
|
println(bst.in_order_traversal())
|
||||||
|
|
||||||
|
bst.insert(KeyVal{ key: 2, val: 34 })
|
||||||
|
bst.insert(KeyVal{ key: -2, val: 203 })
|
||||||
|
|
||||||
|
for elem in bst.in_order_traversal() {
|
||||||
|
println(elem.val)
|
||||||
|
}
|
||||||
|
}
|
|
@ -0,0 +1,296 @@
|
||||||
|
module datatypes
|
||||||
|
|
||||||
|
/// Internal rapresentation of the tree node
|
||||||
|
[heap]
|
||||||
|
struct BSTreeNode<T> {
|
||||||
|
mut:
|
||||||
|
// Mark a node as initialized
|
||||||
|
is_init bool
|
||||||
|
// Value of the node
|
||||||
|
value T
|
||||||
|
// The parent of the node
|
||||||
|
parent &BSTreeNode<T> = 0
|
||||||
|
// The left side with value less than the
|
||||||
|
// value of this node
|
||||||
|
left &BSTreeNode<T> = 0
|
||||||
|
// The right side with value grater than the
|
||||||
|
// value of thiss node
|
||||||
|
right &BSTreeNode<T> = 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// Create new root bst node
|
||||||
|
fn new_root_node<T>(value T) &BSTreeNode<T> {
|
||||||
|
return &BSTreeNode<T>{
|
||||||
|
is_init: true
|
||||||
|
value: value
|
||||||
|
parent: new_none_node<T>(true)
|
||||||
|
left: new_none_node<T>(false)
|
||||||
|
right: new_none_node<T>(false)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// new_node creates a new bst node with a parent reference.
|
||||||
|
fn new_node<T>(parent &BSTreeNode<T>, value T) &BSTreeNode<T> {
|
||||||
|
return &BSTreeNode<T>{
|
||||||
|
is_init: true
|
||||||
|
value: value
|
||||||
|
parent: parent
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// new_none_node creates a dummy node.
|
||||||
|
fn new_none_node<T>(init bool) &BSTreeNode<T> {
|
||||||
|
return &BSTreeNode<T>{
|
||||||
|
is_init: false
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// bind to an actual instance of a node.
|
||||||
|
fn (mut node BSTreeNode<T>) bind(mut to_bind BSTreeNode<T>, left bool) {
|
||||||
|
node.left = to_bind.left
|
||||||
|
node.right = to_bind.right
|
||||||
|
node.value = to_bind.value
|
||||||
|
node.is_init = to_bind.is_init
|
||||||
|
to_bind = new_none_node<T>(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Pure Binary Seach Tree implementation
|
||||||
|
//
|
||||||
|
// Pure V implementation of the Binary Search Tree
|
||||||
|
// Time complexity of main operation O(log N)
|
||||||
|
// Space complexity O(N)
|
||||||
|
pub struct BSTree<T> {
|
||||||
|
mut:
|
||||||
|
root &BSTreeNode<T> = 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// insert give the possibility to insert an element in the BST.
|
||||||
|
pub fn (mut bst BSTree<T>) insert(value T) bool {
|
||||||
|
if bst.is_empty() {
|
||||||
|
bst.root = new_root_node(value)
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return bst.insert_helper(mut bst.root, value)
|
||||||
|
}
|
||||||
|
|
||||||
|
// insert_helper walks the tree and inserts the given node.
|
||||||
|
fn (mut bst BSTree<T>) insert_helper(mut node BSTreeNode<T>, value T) bool {
|
||||||
|
if node.value < value {
|
||||||
|
if node.right != 0 && node.right.is_init {
|
||||||
|
return bst.insert_helper(mut node.right, value)
|
||||||
|
}
|
||||||
|
node.right = new_node(node, value)
|
||||||
|
return true
|
||||||
|
} else if node.value > value {
|
||||||
|
if node.left != 0 && node.left.is_init {
|
||||||
|
return bst.insert_helper(mut node.left, value)
|
||||||
|
}
|
||||||
|
node.left = new_node(node, value)
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
|
||||||
|
// contains checks if an element with a given `value` is inside the BST.
|
||||||
|
pub fn (bst &BSTree<T>) contains(value T) bool {
|
||||||
|
return bst.contains_helper(bst.root, value)
|
||||||
|
}
|
||||||
|
|
||||||
|
// contains_helper is a helper function to walk the tree, and return
|
||||||
|
// the absence or presence of the `value`.
|
||||||
|
fn (bst &BSTree<T>) contains_helper(node &BSTreeNode<T>, value T) bool {
|
||||||
|
if node == 0 || !node.is_init {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if node.value < value {
|
||||||
|
return bst.contains_helper(node.right, value)
|
||||||
|
} else if node.value > value {
|
||||||
|
return bst.contains_helper(node.left, value)
|
||||||
|
}
|
||||||
|
assert node.value == value
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
// remove removes an element with `value` from the BST.
|
||||||
|
pub fn (mut bst BSTree<T>) remove(value T) bool {
|
||||||
|
if bst.root == 0 {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
return bst.remove_helper(mut bst.root, value, false)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn (mut bst BSTree<T>) remove_helper(mut node BSTreeNode<T>, value T, left bool) bool {
|
||||||
|
if !node.is_init {
|
||||||
|
return false
|
||||||
|
}
|
||||||
|
if node.value == value {
|
||||||
|
if node.left != 0 && node.left.is_init {
|
||||||
|
// In order to remove the element we need to bring up as parent the max of the
|
||||||
|
// left sub-tree.
|
||||||
|
mut max_node := bst.get_max_from_right(node.left)
|
||||||
|
node.bind(mut max_node, true)
|
||||||
|
} else if node.right != 0 && node.right.is_init {
|
||||||
|
// Bring up the element with the minimum value in the right sub-tree.
|
||||||
|
mut min_node := bst.get_min_from_left(node.right)
|
||||||
|
node.bind(mut min_node, false)
|
||||||
|
} else {
|
||||||
|
mut parent := node.parent
|
||||||
|
if left {
|
||||||
|
parent.left = new_none_node<T>(false)
|
||||||
|
} else {
|
||||||
|
parent.right = new_none_node<T>(false)
|
||||||
|
}
|
||||||
|
node = new_none_node<T>(false)
|
||||||
|
}
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
|
||||||
|
if node.value < value {
|
||||||
|
return bst.remove_helper(mut node.right, value, false)
|
||||||
|
}
|
||||||
|
return bst.remove_helper(mut node.left, value, true)
|
||||||
|
}
|
||||||
|
|
||||||
|
// get_max_from_right returns the max element of the BST following the right branch.
|
||||||
|
fn (bst &BSTree<T>) get_max_from_right(node &BSTreeNode<T>) &BSTreeNode<T> {
|
||||||
|
right_node := node.right
|
||||||
|
if right_node == 0 || !right_node.is_init {
|
||||||
|
return node
|
||||||
|
}
|
||||||
|
return bst.get_max_from_right(right_node)
|
||||||
|
}
|
||||||
|
|
||||||
|
// get_min_from_left returns the min element of the BST by following the left branch.
|
||||||
|
fn (bst &BSTree<T>) get_min_from_left(node &BSTreeNode<T>) &BSTreeNode<T> {
|
||||||
|
left_node := node.left
|
||||||
|
if left_node == 0 || !left_node.is_init {
|
||||||
|
return node
|
||||||
|
}
|
||||||
|
return bst.get_min_from_left(left_node)
|
||||||
|
}
|
||||||
|
|
||||||
|
// is_empty checks if the BST is empty
|
||||||
|
pub fn (bst &BSTree<T>) is_empty() bool {
|
||||||
|
return bst.root == 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// in_order_traversal traverses the BST in order, and returns the result as an array.
|
||||||
|
pub fn (bst &BSTree<T>) in_order_traversal() []T {
|
||||||
|
mut result := []T{}
|
||||||
|
bst.in_order_traversal_helper(bst.root, mut result)
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// in_order_traversal_helper helps traverse the BST, and accumulates the result in the `result` array.
|
||||||
|
fn (bst &BSTree<T>) in_order_traversal_helper(node &BSTreeNode<T>, mut result []T) {
|
||||||
|
if node == 0 || !node.is_init {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
bst.in_order_traversal_helper(node.left, mut result)
|
||||||
|
result << node.value
|
||||||
|
bst.in_order_traversal_helper(node.right, mut result)
|
||||||
|
}
|
||||||
|
|
||||||
|
// post_order_traversal traverses the BST in post order, and returns the result in an array.
|
||||||
|
pub fn (bst &BSTree<T>) post_order_traversal() []T {
|
||||||
|
mut result := []T{}
|
||||||
|
bst.post_order_traversal_helper(bst.root, mut result)
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// post_order_traversal_helper is a helper function that traverses the BST in post order,
|
||||||
|
// accumulating the result in an array.
|
||||||
|
fn (bst &BSTree<T>) post_order_traversal_helper(node &BSTreeNode<T>, mut result []T) {
|
||||||
|
if node == 0 || !node.is_init {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
|
||||||
|
bst.post_order_traversal_helper(node.left, mut result)
|
||||||
|
bst.post_order_traversal_helper(node.right, mut result)
|
||||||
|
result << node.value
|
||||||
|
}
|
||||||
|
|
||||||
|
// pre_order_traversal traverses the BST in pre order, and returns the result as an array.
|
||||||
|
pub fn (bst &BSTree<T>) pre_order_traversal() []T {
|
||||||
|
mut result := []T{}
|
||||||
|
bst.pre_order_traversal_helper(bst.root, mut result)
|
||||||
|
return result
|
||||||
|
}
|
||||||
|
|
||||||
|
// pre_order_traversal_helper is a helper function to traverse the BST
|
||||||
|
// in pre order and accumulates the results in an array.
|
||||||
|
fn (bst &BSTree<T>) pre_order_traversal_helper(node &BSTreeNode<T>, mut result []T) {
|
||||||
|
if node == 0 || !node.is_init {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
result << node.value
|
||||||
|
bst.pre_order_traversal_helper(node.left, mut result)
|
||||||
|
bst.pre_order_traversal_helper(node.right, mut result)
|
||||||
|
}
|
||||||
|
|
||||||
|
// get_node is a helper method to ge the internal rapresentation of the node with the `value`.
|
||||||
|
fn (bst &BSTree<T>) get_node(node &BSTreeNode<T>, value T) &BSTreeNode<T> {
|
||||||
|
if node == 0 || !node.is_init {
|
||||||
|
return new_none_node<T>(false)
|
||||||
|
}
|
||||||
|
if node.value == value {
|
||||||
|
return node
|
||||||
|
}
|
||||||
|
|
||||||
|
if node.value < value {
|
||||||
|
return bst.get_node(node.right, value)
|
||||||
|
}
|
||||||
|
return bst.get_node(node.left, value)
|
||||||
|
}
|
||||||
|
|
||||||
|
// to_left returns the value of the node to the left of the node with `value` specified if it exists,
|
||||||
|
// otherwise the a false value is returned.
|
||||||
|
//
|
||||||
|
// An example of usage can be the following one
|
||||||
|
//```v
|
||||||
|
// left_value, exist := bst.to_left(10)
|
||||||
|
//```
|
||||||
|
pub fn (bst &BSTree<T>) to_left(value T) ?T {
|
||||||
|
node := bst.get_node(bst.root, value)
|
||||||
|
if !node.is_init {
|
||||||
|
return none
|
||||||
|
}
|
||||||
|
left_node := node.left
|
||||||
|
return left_node.value
|
||||||
|
}
|
||||||
|
|
||||||
|
// to_right return the value of the element to the right of the node with `value` specified, if exist
|
||||||
|
// otherwise, the boolean value is false
|
||||||
|
// An example of usage can be the following one
|
||||||
|
//
|
||||||
|
//```v
|
||||||
|
// left_value, exist := bst.to_right(10)
|
||||||
|
//```
|
||||||
|
pub fn (bst &BSTree<T>) to_right(value T) ?T {
|
||||||
|
node := bst.get_node(bst.root, value)
|
||||||
|
if !node.is_init {
|
||||||
|
return none
|
||||||
|
}
|
||||||
|
right_node := node.right
|
||||||
|
return right_node.value
|
||||||
|
}
|
||||||
|
|
||||||
|
// max return the max element inside the BST.
|
||||||
|
// Time complexity O(N) if the BST is not balanced
|
||||||
|
pub fn (bst &BSTree<T>) max() ?T {
|
||||||
|
max := bst.get_max_from_right(bst.root)
|
||||||
|
if !max.is_init {
|
||||||
|
return none
|
||||||
|
}
|
||||||
|
return max.value
|
||||||
|
}
|
||||||
|
|
||||||
|
// min return the minimum element in the BST.
|
||||||
|
// Time complexity O(N) if the BST is not balanced.
|
||||||
|
pub fn (bst &BSTree<T>) min() ?T {
|
||||||
|
min := bst.get_min_from_left(bst.root)
|
||||||
|
if !min.is_init {
|
||||||
|
return none
|
||||||
|
}
|
||||||
|
return min.value
|
||||||
|
}
|
|
@ -0,0 +1,136 @@
|
||||||
|
module datatypes
|
||||||
|
|
||||||
|
// Make an insert of one element and check if
|
||||||
|
// the bst is able to fin it.
|
||||||
|
fn test_insert_into_bst_one() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10) == true
|
||||||
|
assert bst.contains(10) == true
|
||||||
|
assert bst.contains(20) == false
|
||||||
|
}
|
||||||
|
|
||||||
|
// Make the insert of more element inside the BST
|
||||||
|
// and check if the BST is able to find all the values
|
||||||
|
fn test_insert_into_bst_two() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10)
|
||||||
|
assert bst.insert(20)
|
||||||
|
assert bst.insert(9)
|
||||||
|
|
||||||
|
assert bst.contains(9)
|
||||||
|
assert bst.contains(10)
|
||||||
|
assert bst.contains(20)
|
||||||
|
assert bst.contains(11) == false
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test if the in_order_traversals list return the correct
|
||||||
|
// result array
|
||||||
|
fn test_in_order_bst_visit_one() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10)
|
||||||
|
assert bst.insert(20)
|
||||||
|
assert bst.insert(21)
|
||||||
|
assert bst.insert(1)
|
||||||
|
|
||||||
|
assert bst.in_order_traversal() == [1, 10, 20, 21]
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test if the post_order_bst_visit return the correct
|
||||||
|
// result array
|
||||||
|
fn test_post_order_bst_visit_one() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10)
|
||||||
|
assert bst.insert(20)
|
||||||
|
assert bst.insert(21)
|
||||||
|
assert bst.insert(1)
|
||||||
|
|
||||||
|
assert bst.post_order_traversal() == [1, 21, 20, 10]
|
||||||
|
}
|
||||||
|
|
||||||
|
// Test if the pre_order_traversal return the correct result array
|
||||||
|
fn test_pre_order_bst_visit_one() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10)
|
||||||
|
assert bst.insert(20)
|
||||||
|
assert bst.insert(21)
|
||||||
|
assert bst.insert(1)
|
||||||
|
|
||||||
|
assert bst.pre_order_traversal() == [10, 1, 20, 21]
|
||||||
|
}
|
||||||
|
|
||||||
|
// After many insert check if we are abe to get the correct
|
||||||
|
// right and left value of the root.
|
||||||
|
fn test_get_left_root() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10)
|
||||||
|
assert bst.insert(20)
|
||||||
|
assert bst.insert(21)
|
||||||
|
assert bst.insert(1)
|
||||||
|
|
||||||
|
left_val := bst.to_left(10) or { -1 }
|
||||||
|
assert left_val == 1
|
||||||
|
|
||||||
|
right_val := bst.to_right(10) or { -1 }
|
||||||
|
assert right_val == 20
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check if BST panic if we call some operation on an empty BST.
|
||||||
|
fn test_get_left_on_empty_bst() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
|
||||||
|
left_val := bst.to_left(10) or { -1 }
|
||||||
|
assert left_val == -1
|
||||||
|
|
||||||
|
right_val := bst.to_right(10) or { -1 }
|
||||||
|
assert right_val == -1
|
||||||
|
}
|
||||||
|
|
||||||
|
// Check the remove operation if it is able to remove
|
||||||
|
// all elements required, and mantains the BST propriety.
|
||||||
|
fn test_remove_from_bst_one() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10)
|
||||||
|
assert bst.insert(20)
|
||||||
|
assert bst.insert(21)
|
||||||
|
assert bst.insert(1)
|
||||||
|
assert bst.in_order_traversal() == [1, 10, 20, 21]
|
||||||
|
assert bst.remove(21)
|
||||||
|
|
||||||
|
assert bst.in_order_traversal() == [1, 10, 20]
|
||||||
|
}
|
||||||
|
|
||||||
|
// Another test n the remove BST, this remove an intermidia node
|
||||||
|
// that it is a triky operation.
|
||||||
|
fn test_remove_from_bst_two() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10)
|
||||||
|
assert bst.insert(20)
|
||||||
|
assert bst.insert(21)
|
||||||
|
assert bst.insert(1)
|
||||||
|
assert bst.in_order_traversal() == [1, 10, 20, 21]
|
||||||
|
assert bst.remove(20)
|
||||||
|
|
||||||
|
assert bst.in_order_traversal() == [1, 10, 21]
|
||||||
|
}
|
||||||
|
|
||||||
|
// check if we are able to get the max from the BST.
|
||||||
|
fn test_get_max_in_bst() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10)
|
||||||
|
assert bst.insert(20)
|
||||||
|
assert bst.insert(21)
|
||||||
|
assert bst.insert(1)
|
||||||
|
max := bst.max() or { -1 }
|
||||||
|
assert max == 21
|
||||||
|
}
|
||||||
|
|
||||||
|
// check if we are able to get the min from the BST.
|
||||||
|
fn test_get_min_in_bst() {
|
||||||
|
mut bst := BSTree<int>{}
|
||||||
|
assert bst.insert(10)
|
||||||
|
assert bst.insert(20)
|
||||||
|
assert bst.insert(21)
|
||||||
|
assert bst.insert(1)
|
||||||
|
min := bst.min() or { -1 }
|
||||||
|
assert min == 1
|
||||||
|
}
|
Loading…
Reference in New Issue