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