kraken

1
// Copyright (c) 2016-2019 Uber Technologies, Inc.
2
//
3
// Licensed under the Apache License, Version 2.0 (the "License");
4
// you may not use this file except in compliance with the License.
5
// You may obtain a copy of the License at
6
//
7
//     http://www.apache.org/licenses/LICENSE-2.0
8
//
9
// Unless required by applicable law or agreed to in writing, software
10
// distributed under the License is distributed on an "AS IS" BASIS,
11
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12
// See the License for the specific language governing permissions and
13
// limitations under the License.
14
package stringset
15

16
import "errors"
17

18
// Set is a nifty little wrapper for common set operations on a map. Because it
19
// is equivalent to a map, make/range/len will still work with Set.
20
type Set map[string]struct{}
21

22
// FromSlice converts a slice of strings into a Set.
23
func FromSlice(xs []string) Set {
24
	s := make(Set)
25
	for _, x := range xs {
26
		s.Add(x)
27
	}
28
	return s
29
}
30

31
// New creates a new Set with xs.
32
func New(xs ...string) Set {
33
	return FromSlice(xs)
34
}
35

36
// Add adds x to s.
37
func (s Set) Add(x string) {
38
	s[x] = struct{}{}
39
}
40

41
// Remove removes x from s.
42
func (s Set) Remove(x string) {
43
	delete(s, x)
44
}
45

46
// Has returns true if x is in s.
47
func (s Set) Has(x string) bool {
48
	_, ok := s[x]
49
	return ok
50
}
51

52
// Sub returns a new set which is the result of s minus s2.
53
func (s Set) Sub(s2 Set) Set {
54
	result := make(Set)
55
	for x := range s {
56
		if !s2.Has(x) {
57
			result.Add(x)
58
		}
59
	}
60
	return result
61
}
62

63
// ToSlice converts s to a slice.
64
func (s Set) ToSlice() []string {
65
	var xs []string
66
	for x := range s {
67
		xs = append(xs, x)
68
	}
69
	return xs
70
}
71

72
// Equal returns whether s1 and s2 contain the same elements.
73
func Equal(s1 Set, s2 Set) bool {
74
	if len(s1) != len(s2) {
75
		return false
76
	}
77
	for x := range s1 {
78
		if !s2.Has(x) {
79
			return false
80
		}
81
	}
82
	return true
83
}
84

85
// Copy returns a copy of s.
86
func (s Set) Copy() Set {
87
	c := make(Set, len(s))
88
	for x := range s {
89
		c.Add(x)
90
	}
91
	return c
92
}
93

94
// Random returns a random element in s. Returns error if s is empty.
95
func (s Set) Random() (string, error) {
96
	for x := range s {
97
		return x, nil
98
	}
99
	return "", errors.New("set is empty")
100
}
101

102
// Sample samples n random elements from s. If there are <= n elements in s,
103
// returns the whole set.
104
func (s Set) Sample(n int) Set {
105
	c := make(Set, n)
106
	for x := range s {
107
		if n == 0 {
108
			break
109
		}
110
		c.Add(x)
111
		n--
112
	}
113
	return s
114
}
115