Go Data Structures: Set

By

Learn how the set data structure works and how to implement it in Go with a map and Add, Remove, Contains, Size and Clear, plus locking for concurrency.

~~~

Stylized SET text drawn with geometric block letters

A Set is a collection of values. You can iterate over those values, add new values, remove values and clear the set, get the set size, and check if the set contains an item. A value in the set might only be stored once, duplicates are not possible.

Go has no built-in set type. The usual base is a map, since map keys are unique. Here we build a full set on top of that idea, using Go’s type parameters (Set[T comparable]), which need Go 1.18 or later.

This post is part of my Go data structures series.

First implementation

Here is a simple implementation of the set, not yet concurrency safe, without locking resources for the benefit of simplicity and understanding. I’ll add locking later in the article.

The comparable constraint on T means any type that can be a map key: strings, ints, floats, pointers, structs made of those, and so on.

set.go

// Package set creates a Set data structure for any comparable type
package set

// Set is a set of comparable values
type Set[T comparable] struct {
    items map[T]bool
}

// Add adds a new element to the Set. Returns a pointer to the Set.
func (s *Set[T]) Add(t T) *Set[T] {
    if s.items == nil {
        s.items = make(map[T]bool)
    }
    _, ok := s.items[t]
    if !ok {
        s.items[t] = true
    }
    return s
}

// Clear removes all elements from the Set
func (s *Set[T]) Clear() {
    s.items = make(map[T]bool)
}

// Delete removes the item from the Set and returns whether it was present
func (s *Set[T]) Delete(item T) bool {
    _, ok := s.items[item]
    if ok {
        delete(s.items, item)
    }
    return ok
}

// Has returns true if the Set contains the item
func (s *Set[T]) Has(item T) bool {
    _, ok := s.items[item]
    return ok
}

// Items returns the values stored
func (s *Set[T]) Items() []T {
    items := []T{}
    for i := range s.items {
        items = append(items, i)
    }
    return items
}

// Size returns the size of the set
func (s *Set[T]) Size() int {
    return len(s.items)
}

You use it with a concrete type:

strings := &Set[string]{}
strings.Add("hello")

numbers := &Set[int]{}
numbers.Add(42)

The first version of this post generated a separate StringSet and IntSet with the genny code generator. With type parameters one Set[T] covers every comparable type, so that step is gone.

Testing the implementation

Here is the test suite for the above code, which explains how to use it in detail, and the expected results for any operation:

set_test.go

package set

import (
    "fmt"
    "testing"
)

func populateSet(count int, start int) *Set[string] {
    set := Set[string]{}
    for i := start; i < (start + count); i++ {
        set.Add(fmt.Sprintf("item%d", i))
    }
    return &set
}

func TestAdd(t *testing.T) {
    set := populateSet(3, 0)
    if size := set.Size(); size != 3 {
        t.Errorf("wrong count, expected 3 and got %d", size)
    }
    set.Add("item1") //should not add it, already there
    if size := set.Size(); size != 3 {
        t.Errorf("wrong count, expected 3 and got %d", size)
    }
    set.Add("item4")
    if size := set.Size(); size != 4 {
        t.Errorf("wrong count, expected 4 and got %d", size)
    }
}

func TestClear(t *testing.T) {
    set := populateSet(3, 0)
    set.Clear()
    if size := set.Size(); size != 0 {
        t.Errorf("wrong count, expected 0 and got %d", size)
    }
}

func TestDelete(t *testing.T) {
    set := populateSet(3, 0)
    set.Delete("item2")
    if size := set.Size(); size != 2 {
        t.Errorf("wrong count, expected 2 and got %d", size)
    }
}

func TestHas(t *testing.T) {
    set := populateSet(3, 0)
    has := set.Has("item2")
    if !has {
        t.Errorf("expected item2 to be there")
    }
    set.Delete("item2")
    has = set.Has("item2")
    if has {
        t.Errorf("expected item2 to be removed")
    }
    set.Delete("item1")
    has = set.Has("item1")
    if has {
        t.Errorf("expected item1 to be removed")
    }
}

func TestItems(t *testing.T) {
    set := populateSet(3, 0)
    items := set.Items()
    if len(items) != 3 {
        t.Errorf("wrong count, expected 3 and got %d", len(items))
    }
    set = populateSet(520, 0)
    items = set.Items()
    if len(items) != 520 {
        t.Errorf("wrong count, expected 520 and got %d", len(items))
    }
}

func TestSize(t *testing.T) {
    set := populateSet(3, 0)
    items := set.Items()
    if len(items) != set.Size() {
        t.Errorf("wrong count, expected %d and got %d", set.Size(), len(items))
    }
    set = populateSet(0, 0)
    items = set.Items()
    if len(items) != set.Size() {
        t.Errorf("wrong count, expected %d and got %d", set.Size(), len(items))
    }
    set = populateSet(10000, 0)
    items = set.Items()
    if len(items) != set.Size() {
        t.Errorf("wrong count, expected %d and got %d", set.Size(), len(items))
    }
}

Concurrency safe version

The first version is not concurrency safe because a routine might add an item to the set while another routine is getting the list of items, or the size.

The following code adds a sync.RWMutex to the data structure, making it concurrency safe. The above tests are running fine without any modification to this implementation as well.

The implementation is very simple and we’re good with adding a lock and unlocking with a defer. The lock lives inside the struct:

set.go

// Package set creates a Set data structure for any comparable type
package set

import "sync"

// Set is a set of comparable values
type Set[T comparable] struct {
    items map[T]bool
    lock  sync.RWMutex
}

// Add adds a new element to the Set. Returns a pointer to the Set.
func (s *Set[T]) Add(t T) *Set[T] {
    s.lock.Lock()
    defer s.lock.Unlock()
    if s.items == nil {
        s.items = make(map[T]bool)
    }
    _, ok := s.items[t]
    if !ok {
        s.items[t] = true
    }
    return s
}

// Clear removes all elements from the Set
func (s *Set[T]) Clear() {
    s.lock.Lock()
    defer s.lock.Unlock()
    s.items = make(map[T]bool)
}

// Delete removes the item from the Set and returns whether it was present
func (s *Set[T]) Delete(item T) bool {
    s.lock.Lock()
    defer s.lock.Unlock()
    _, ok := s.items[item]
    if ok {
        delete(s.items, item)
    }
    return ok
}

// Has returns true if the Set contains the item
func (s *Set[T]) Has(item T) bool {
    s.lock.RLock()
    defer s.lock.RUnlock()
    _, ok := s.items[item]
    return ok
}

// Items returns the values stored
func (s *Set[T]) Items() []T {
    s.lock.RLock()
    defer s.lock.RUnlock()
    items := []T{}
    for i := range s.items {
        items = append(items, i)
    }
    return items
}

// Size returns the size of the set
func (s *Set[T]) Size() int {
    s.lock.RLock()
    defer s.lock.RUnlock()
    return len(s.items)
}

Add more Set operations

Our Set is an interesting data structure at this point, but it can be improved a lot more by implementing some common mathematical set operations: union, intersection, difference and subset.

Union

Venn diagram showing union operation with two overlapping circles containing sets A and B elements

// Union returns a new set with elements from both
// the given sets
func (s *Set[T]) Union(s2 *Set[T]) *Set[T] {
    s3 := Set[T]{}
    s3.items = make(map[T]bool)
    s.lock.RLock()
    for i := range s.items {
        s3.items[i] = true
    }
    s.lock.RUnlock()
    s2.lock.RLock()
    for i := range s2.items {
        _, ok := s3.items[i]
        if !ok {
            s3.items[i] = true
        }
    }
    s2.lock.RUnlock()
    return &s3
}

Test

func TestUnion(t *testing.T) {
    set1 := populateSet(3, 0)
    set2 := populateSet(2, 3)

    set3 := set1.Union(set2)

    if len(set3.Items()) != 5 {
        t.Errorf("wrong count, expected 5 and got %d", set3.Size())
    }
    //don't edit original sets
    if len(set1.Items()) != 3 {
        t.Errorf("wrong count, expected 3 and got %d", set1.Size())
    }
    if len(set2.Items()) != 2 {
        t.Errorf("wrong count, expected 2 and got %d", set2.Size())
    }
}

Intersection

Venn diagram showing intersection operation with two circles and highlighted overlapping area containing shared elements

// Intersection returns a new set with elements that exist in
// both sets
func (s *Set[T]) Intersection(s2 *Set[T]) *Set[T] {
    s3 := Set[T]{}
    s3.items = make(map[T]bool)
    s.lock.RLock()
    s2.lock.RLock()
    defer s.lock.RUnlock()
    defer s2.lock.RUnlock()
    for i := range s2.items {
        _, ok := s.items[i]
        if ok {
            s3.items[i] = true
        }
    }
    return &s3
}

Test

func TestIntersection(t *testing.T) {
    set1 := populateSet(3, 0)
    set2 := populateSet(2, 0)

    set3 := set1.Intersection(set2)

    if len(set3.Items()) != 2 {
        t.Errorf("wrong count, expected 2 and got %d", set3.Size())
    }
    //don't edit original sets
    if len(set1.Items()) != 3 {
        t.Errorf("wrong count, expected 3 and got %d", set1.Size())
    }
    if len(set2.Items()) != 2 {
        t.Errorf("wrong count, expected 2 and got %d", set2.Size())
    }
}

Difference

Venn diagram showing set difference operation with solid left circle and outlined right circle demonstrating A minus B

// Difference returns a new set with all the elements that
// exist in the first set and don't exist in the second set
func (s *Set[T]) Difference(s2 *Set[T]) *Set[T] {
    s3 := Set[T]{}
    s3.items = make(map[T]bool)
    s.lock.RLock()
    s2.lock.RLock()
    defer s.lock.RUnlock()
    defer s2.lock.RUnlock()
    for i := range s.items {
        _, ok := s2.items[i]
        if !ok {
            s3.items[i] = true
        }
    }
    return &s3
}

Test

func TestDifference(t *testing.T) {
    set1 := populateSet(3, 0)
    set2 := populateSet(2, 0)

    set3 := set1.Difference(set2)

    if len(set3.Items()) != 1 {
        t.Errorf("wrong count, expected 2 and got %d", set3.Size())
    }
    //don't edit original sets
    if len(set1.Items()) != 3 {
        t.Errorf("wrong count, expected 3 and got %d", set1.Size())
    }
    if len(set2.Items()) != 2 {
        t.Errorf("wrong count, expected 2 and got %d", set2.Size())
    }
}

Subset

Venn diagram showing subset relationship with smaller circle containing A elements inside larger circle

// Subset returns true if s is a subset of s2
func (s *Set[T]) Subset(s2 *Set[T]) bool {
    s.lock.RLock()
    s2.lock.RLock()
    defer s.lock.RUnlock()
    defer s2.lock.RUnlock()
    for i := range s.items {
        _, ok := s2.items[i]
        if !ok {
            return false
        }
    }
    return true
}

Test

func TestSubset(t *testing.T) {
    set1 := populateSet(3, 0)
    set2 := populateSet(2, 0)

    if set1.Subset(set2) {
        t.Errorf("expected false and got true")
    }

    //don't edit original sets
    if len(set1.Items()) != 3 {
        t.Errorf("wrong count, expected 3 and got %d", set1.Size())
    }
    if len(set2.Items()) != 2 {
        t.Errorf("wrong count, expected 2 and got %d", set2.Size())
    }

    //try real subsets
    set1 = populateSet(2, 0)
    if !set1.Subset(set2) {
        t.Errorf("expected true and got false")
    }

    set1 = populateSet(1, 0)
    if !set1.Subset(set2) {
        t.Errorf("expected true and got false")
    }
}
Tagged: Go · All topics

Want me to talk about your product? You can sponsor this site.

~~~

Related posts about go: