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