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