gitee.com/quant1x/gox@v1.21.2/util/examples/redblacktreeextended/redblacktreeextended.go (about)

     1  // Copyright (c) 2015, Emir Pasic. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package redblacktreeextended
     6  
     7  import (
     8  	"fmt"
     9  	rbt "gitee.com/quant1x/gox/util/redblacktree"
    10  )
    11  
    12  // RedBlackTreeExtended to demonstrate how to extend a RedBlackTree to include new functions
    13  type RedBlackTreeExtended struct {
    14  	*rbt.Tree
    15  }
    16  
    17  // GetMin gets the min value and flag if found
    18  func (tree *RedBlackTreeExtended) GetMin() (value interface{}, found bool) {
    19  	node, found := tree.getMinFromNode(tree.Root)
    20  	if node != nil {
    21  		return node.Value, found
    22  	}
    23  	return nil, false
    24  }
    25  
    26  // GetMax gets the max value and flag if found
    27  func (tree *RedBlackTreeExtended) GetMax() (value interface{}, found bool) {
    28  	node, found := tree.getMaxFromNode(tree.Root)
    29  	if node != nil {
    30  		return node.Value, found
    31  	}
    32  	return nil, false
    33  }
    34  
    35  // RemoveMin removes the min value and flag if found
    36  func (tree *RedBlackTreeExtended) RemoveMin() (value interface{}, deleted bool) {
    37  	node, found := tree.getMinFromNode(tree.Root)
    38  	if found {
    39  		tree.Remove(node.Key)
    40  		return node.Value, found
    41  	}
    42  	return nil, false
    43  }
    44  
    45  // RemoveMax removes the max value and flag if found
    46  func (tree *RedBlackTreeExtended) RemoveMax() (value interface{}, deleted bool) {
    47  	node, found := tree.getMaxFromNode(tree.Root)
    48  	if found {
    49  		tree.Remove(node.Key)
    50  		return node.Value, found
    51  	}
    52  	return nil, false
    53  }
    54  
    55  func (tree *RedBlackTreeExtended) getMinFromNode(node *rbt.Node) (foundNode *rbt.Node, found bool) {
    56  	if node == nil {
    57  		return nil, false
    58  	}
    59  	if node.Left == nil {
    60  		return node, true
    61  	}
    62  	return tree.getMinFromNode(node.Left)
    63  }
    64  
    65  func (tree *RedBlackTreeExtended) getMaxFromNode(node *rbt.Node) (foundNode *rbt.Node, found bool) {
    66  	if node == nil {
    67  		return nil, false
    68  	}
    69  	if node.Right == nil {
    70  		return node, true
    71  	}
    72  	return tree.getMaxFromNode(node.Right)
    73  }
    74  
    75  func print(tree *RedBlackTreeExtended) {
    76  	max, _ := tree.GetMax()
    77  	min, _ := tree.GetMin()
    78  	fmt.Printf("Value for max key: %v \n", max)
    79  	fmt.Printf("Value for min key: %v \n", min)
    80  	fmt.Println(tree)
    81  }
    82  
    83  // RedBlackTreeExtendedExample main method on how to use the custom red-black tree above
    84  func main() {
    85  	tree := RedBlackTreeExtended{rbt.NewWithIntComparator()}
    86  
    87  	tree.Put(1, "a") // 1->x (in order)
    88  	tree.Put(2, "b") // 1->x, 2->b (in order)
    89  	tree.Put(3, "c") // 1->x, 2->b, 3->c (in order)
    90  	tree.Put(4, "d") // 1->x, 2->b, 3->c, 4->d (in order)
    91  	tree.Put(5, "e") // 1->x, 2->b, 3->c, 4->d, 5->e (in order)
    92  
    93  	print(&tree)
    94  	// Value for max key: e
    95  	// Value for min key: a
    96  	// RedBlackTree
    97  	// │       ┌── 5
    98  	// │   ┌── 4
    99  	// │   │   └── 3
   100  	// └── 2
   101  	//     └── 1
   102  
   103  	tree.RemoveMin() // 2->b, 3->c, 4->d, 5->e (in order)
   104  	tree.RemoveMax() // 2->b, 3->c, 4->d (in order)
   105  	tree.RemoveMin() // 3->c, 4->d (in order)
   106  	tree.RemoveMax() // 3->c (in order)
   107  
   108  	print(&tree)
   109  	// Value for max key: c
   110  	// Value for min key: c
   111  	// RedBlackTree
   112  	// └── 3
   113  }