github.com/Nigel2392/go-datastructures@v1.1.5/binarytree/bst.go (about)

     1  package binarytree
     2  
     3  import (
     4  	"fmt"
     5  	"math"
     6  	"strings"
     7  
     8  	"github.com/Nigel2392/go-datastructures"
     9  	"golang.org/x/exp/slices"
    10  )
    11  
    12  // A binary search tree implementation.
    13  type BST[T datastructures.Ordered] struct {
    14  	root   *BSTNode[T]
    15  	len    int
    16  	height int
    17  }
    18  
    19  // Return the binary search tree as a string.
    20  func (t *BST[T]) String() string {
    21  	if t.root == nil {
    22  		return ""
    23  	}
    24  
    25  	height := t.root.getHeight()
    26  	BSTNodes := make([][]string, height)
    27  
    28  	fillBSTNodes(BSTNodes, t.root, 0)
    29  
    30  	var b strings.Builder
    31  	padding := int(math.Pow(2, float64(height)) - 1)
    32  
    33  	for i, level := range BSTNodes {
    34  		if i == 0 {
    35  			paddingStr := strings.Repeat(" ", (padding/2)+1)
    36  			b.WriteString(paddingStr)
    37  		} else {
    38  			paddingStr := strings.Repeat(" ", padding/2)
    39  			b.WriteString(paddingStr)
    40  		}
    41  
    42  		for j, BSTNode := range level {
    43  			b.WriteString(BSTNode)
    44  			if j != len(level)-1 {
    45  				b.WriteString(strings.Repeat(" ", padding))
    46  			}
    47  		}
    48  
    49  		padding /= 2
    50  		b.WriteString("\n")
    51  	}
    52  
    53  	return b.String()
    54  }
    55  
    56  // Initialize a new binary search tree with the given initial value.
    57  func NewBST[T datastructures.Ordered](initial T) *BST[T] {
    58  	return &BST[T]{
    59  		root: &BSTNode[T]{value: initial}}
    60  }
    61  
    62  // Insert a new value into the binary search tree.
    63  func (t *BST[T]) Insert(value T) (inserted bool) {
    64  	if t.root == nil {
    65  		t.root = &BSTNode[T]{value: value}
    66  		t.len++
    67  		return true
    68  	}
    69  	inserted = t.root.insert(value)
    70  	if inserted {
    71  		t.len++
    72  	}
    73  	return inserted
    74  }
    75  
    76  // Search for a value in the binary search tree.
    77  func (t *BST[T]) Search(value T) (v T, ok bool) {
    78  	if t.root == nil {
    79  		return
    80  	}
    81  	return t.root.search(value)
    82  }
    83  
    84  // Delete a value from the binary search tree.
    85  func (t *BST[T]) Delete(value T) (deleted bool) {
    86  	if t.root == nil {
    87  		return false
    88  	}
    89  	t.root, deleted = t.root.delete(value)
    90  	if deleted {
    91  		t.len--
    92  	}
    93  	return deleted
    94  }
    95  
    96  // Delete a value from the binary search tree if the predicate returns true.
    97  func (t *BST[T]) DeleteIf(predicate func(T) bool) (deleted int) {
    98  	if t.root == nil {
    99  		return 0
   100  	}
   101  	t.root, deleted = t.root.deleteIf(predicate)
   102  	t.len -= int(deleted)
   103  	return deleted
   104  }
   105  
   106  // Traverse the binary search tree in order.
   107  func (t *BST[T]) Traverse(f func(T)) {
   108  	if t.root == nil {
   109  		return
   110  	}
   111  	t.root.traverse(f)
   112  }
   113  
   114  // Return the number of values in the binary search tree.
   115  func (t *BST[T]) Len() int {
   116  	return t.len
   117  }
   118  
   119  // Return the height of the binary search tree.
   120  func (t *BST[T]) Height() int {
   121  	if t.root == nil {
   122  		return 0
   123  	}
   124  	return t.root.getHeight()
   125  }
   126  
   127  func fillBSTNodes[T datastructures.Ordered](BSTNodes [][]string, n *BSTNode[T], depth int) {
   128  	if n == nil {
   129  		return
   130  	}
   131  
   132  	BSTNodes[depth] = append(BSTNodes[depth], fmt.Sprintf("%v", n.value))
   133  	fillBSTNodes(BSTNodes, n.left, depth+1)
   134  	fillBSTNodes(BSTNodes, n.right, depth+1)
   135  }
   136  
   137  func SliceToBST[T datastructures.Ordered](items []T, sorted bool) *BST[T] {
   138  	if !sorted {
   139  		slices.Sort(items)
   140  	}
   141  	var bst BST[T]
   142  	bst.root = constructBSTFromSortedSlice(items, 0, len(items))
   143  	return &bst
   144  }
   145  
   146  func constructBSTFromSortedSlice[T datastructures.Ordered](items []T, start, end int) *BSTNode[T] {
   147  	if start == end {
   148  		return nil
   149  	}
   150  	mid := start + (end-start)/2
   151  	return &BSTNode[T]{
   152  		value: items[mid],
   153  		left:  constructBSTFromSortedSlice(items, start, mid),
   154  		right: constructBSTFromSortedSlice(items, mid+1, end),
   155  	}
   156  }