gitee.com/quant1x/gox@v1.21.2/util/treeset/treeset.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 treeset implements a tree backed by a red-black tree.
     6  //
     7  // Structure is not thread safe.
     8  //
     9  // Reference: http://en.wikipedia.org/wiki/Set_%28abstract_data_type%29
    10  package treeset
    11  
    12  import (
    13  	"fmt"
    14  	"gitee.com/quant1x/gox/util/internal"
    15  	rbt "gitee.com/quant1x/gox/util/redblacktree"
    16  	"strings"
    17  )
    18  
    19  func assertSetImplementation() {
    20  	var _ internal.Set = (*Set)(nil)
    21  }
    22  
    23  // Set holds elements in a red-black tree
    24  type Set struct {
    25  	tree *rbt.Tree
    26  }
    27  
    28  var itemExists = struct{}{}
    29  
    30  // NewWith instantiates a new empty set with the custom comparator.
    31  func NewWith(comparator internal.Comparator, values ...interface{}) *Set {
    32  	set := &Set{tree: rbt.NewWith(comparator)}
    33  	if len(values) > 0 {
    34  		set.Add(values...)
    35  	}
    36  	return set
    37  }
    38  
    39  // NewWithIntComparator instantiates a new empty set with the IntComparator, i.e. keys are of type int.
    40  func NewWithIntComparator(values ...interface{}) *Set {
    41  	set := &Set{tree: rbt.NewWithIntComparator()}
    42  	if len(values) > 0 {
    43  		set.Add(values...)
    44  	}
    45  	return set
    46  }
    47  
    48  // NewWithStringComparator instantiates a new empty set with the StringComparator, i.e. keys are of type string.
    49  func NewWithStringComparator(values ...interface{}) *Set {
    50  	set := &Set{tree: rbt.NewWithStringComparator()}
    51  	if len(values) > 0 {
    52  		set.Add(values...)
    53  	}
    54  	return set
    55  }
    56  
    57  // Add adds the items (one or more) to the set.
    58  func (set *Set) Add(items ...interface{}) {
    59  	for _, item := range items {
    60  		set.tree.Put(item, itemExists)
    61  	}
    62  }
    63  
    64  // Remove removes the items (one or more) from the set.
    65  func (set *Set) Remove(items ...interface{}) {
    66  	for _, item := range items {
    67  		set.tree.Remove(item)
    68  	}
    69  }
    70  
    71  // Contains checks weather items (one or more) are present in the set.
    72  // All items have to be present in the set for the method to return true.
    73  // Returns true if no arguments are passed at all, i.e. set is always superset of empty set.
    74  func (set *Set) Contains(items ...interface{}) bool {
    75  	for _, item := range items {
    76  		if _, contains := set.tree.Get(item); !contains {
    77  			return false
    78  		}
    79  	}
    80  	return true
    81  }
    82  
    83  // Empty returns true if set does not contain any elements.
    84  func (set *Set) Empty() bool {
    85  	return set.tree.Size() == 0
    86  }
    87  
    88  // Size returns number of elements within the set.
    89  func (set *Set) Size() int {
    90  	return set.tree.Size()
    91  }
    92  
    93  // Clear clears all values in the set.
    94  func (set *Set) Clear() {
    95  	set.tree.Clear()
    96  }
    97  
    98  // Values returns all items in the set.
    99  func (set *Set) Values() []interface{} {
   100  	return set.tree.Keys()
   101  }
   102  
   103  // String returns a string representation of container
   104  func (set *Set) String() string {
   105  	str := "TreeSet\n"
   106  	items := []string{}
   107  	for _, v := range set.tree.Keys() {
   108  		items = append(items, fmt.Sprintf("%v", v))
   109  	}
   110  	str += strings.Join(items, ", ")
   111  	return str
   112  }