exercism/go/custom-set/stringset.go

package stringset

import (
	"errors"
	"fmt"
	"strings"
)

const testVersion = 3

// Set is a slice of strings that you can do set operations on.
// I decided that I wanted to implement a binary tree for the storage
type Set struct {
	top SetValue
}

type SetValue struct {
	value string
	left  *SetValue
	right *SetValue
}

// New returns an empty Set
func New() Set {
	s := new(Set)
	s.top = SetValue{}
	return *s
}

// NewFromSlice takes a slice of strings and returns a Set
func NewFromSlice(s []string) Set {
	ret := New()
	for i := range s {
		ret.Add(s[i])
	}
	ret.balance()
	return ret
}

// Add adds a value to the set
func (s Set) Add(v string) {
	if s.top.value == "" {
		s.top.value = v
		return
	}
	s.top.Add(v)
}

func (s *Set) PrettyPrint() {
	s.pp(&s.top, 0)
}

func (s *Set) pp(n *SetValue, indent int) {
	if n != nil {
		if n.left != nil {
			s.pp(n.left, indent+4)
		}
		if n.right != nil {
			s.pp(n.right, indent+4)
		}
		for ; indent > 0; indent-- {
			fmt.Print(" ")
		}
		fmt.Println(n.value)
	}
}

// balance balances the binary tree
func (s *Set) balance() {

}

func (s *Set) find(v string) *SetValue {
	if s.top.value != "" {
		sv, _ := s.top.find(v)
		return sv
	}
	return nil
}

// findWithParent finds a node with a child value v
// a return of nil, nil means it's the top node
func (s *Set) findParent(v string) (*SetValue, error) {
	if s.top.value == v {
		// no parent, it's the top.
		return nil, nil
	} else {
		return s.top.findParent(v)
	}
	return nil, errors.New("Empty Set")
}

func (s *Set) findHome(v *SetValue) {
	if s.top.value == "" {
		s.top.value = v.value
		s.top.left = v.left
		s.top.right = v.right
		return
	}
	s.top.findHome(v)
}

// Delete removes the given value from the set
func (s Set) Delete(v string) {
	if sv, err := s.findParent(v); err == nil {
		var cmp int
		var delNode, repNode, orphan *SetValue
		if sv == nil {
			// Deleting 'top'
			delNode = &s.top
			if delNode.left != nil {
				repNode = delNode.left
				orphan = delNode.right
			} else if delNode.right != nil {
				repNode = delNode.right
			}
			if repNode == nil {
				// No node to replace it with, we're done
				return
			}
			s.top = *repNode
		} else {
			cmp = strings.Compare(v, sv.value)
			if cmp < 0 && sv.left != nil {
				fmt.Println("    Left: " + sv.left.value)
				// It's the left node
				delNode = sv.left
			} else if cmp > 0 && sv.right != nil {
				fmt.Println("    Right: " + sv.right.value)
				// It's the right node
				delNode = sv.right
			}
			if delNode == nil {
				return
			}

			if delNode.left != nil {
				repNode = delNode.left
				orphan = delNode.right
			} else if delNode.right != nil {
				repNode = delNode.right
			}
			if repNode == nil {
				// No replacement node, we're done
				return
			}
			if cmp < 0 {
				sv.left = repNode
			} else if cmp > 0 {
				sv.right = repNode
			}
		}
		// If we have an orphaned branch, find it a home
		if orphan != nil {
			s.findHome(orphan)
		}
	}
}

// Has returns if the set contains the given value.
func (s *Set) Has(v string) bool {
	return s.find(v) != nil
}

// IsEmpty returns whether the set is empty or not.
func (s *Set) IsEmpty() bool {
	return s.top.value == ""
}

// Len returns the number of values in the set
func (s *Set) Len() int {
	if !s.IsEmpty() {
		return s.top.Len()
	}
	return 0
}

// Slice returns a string slice of the set
func (s Set) Slice() []string {
	if !s.IsEmpty() {
		return s.top.Slice()
	}
	return []string{}
}

// String converts the set to a string
func (s Set) String() string {
	ret := "{"
	if s.top.value != "" {
		ret += s.top.String()
	}
	ret += "}"
	return ret
}

// find looks for a node with value val, it either returns the node
// or an error stating it couldn't find it.
func (sv *SetValue) find(val string) (*SetValue, error) {
	if sv.value == val {
		return sv, nil
	}
	cmp := strings.Compare(val, sv.value)
	if cmp < 0 && sv.left != nil {
		return sv.left.find(val)
	}
	if cmp > 0 && sv.right != nil {
		return sv.right.find(val)
	}
	return nil, errors.New("Value not found")
}

// findParent looks for the parent of the node with value val
// If nil, nil is returned, it _is_ this node.
func (sv *SetValue) findParent(val string) (*SetValue, error) {
	if sv.value == val {
		// This should only trigger if this is the top node of the tree
		return nil, nil
	}
	cmp := strings.Compare(val, sv.value)
	if cmp < 0 && sv.left != nil {
		if sv.left.value == val {
			return sv, nil
		}
		return sv.left.findParent(val)
	}
	if cmp > 0 && sv.right != nil {
		if sv.right.value == val {
			return sv, nil
		}
		return sv.right.findParent(val)
	}
	return nil, errors.New("Value not found")
}

func (sv *SetValue) findHome(v *SetValue) {
	cmp := strings.Compare(v.value, sv.value)
	if cmp < 0 {
		if sv.left == nil {
			sv.left = v
		} else {
			sv.left.findHome(v)
		}
	} else if cmp > 0 {
		if sv.right == nil {
			sv.right = v
		} else {
			sv.right.findHome(v)
		}
	} else {
		// Discard the top node, find homes for it's children
		if v.left != nil {
			sv.findHome(v.left)
		}
		if v.right != nil {
			sv.findHome(v.right)
		}
	}
}

func (sv *SetValue) Add(v string) {
	cmp := strings.Compare(v, sv.value)
	if cmp < 0 {
		if sv.left == nil {
			sv.left = &SetValue{value: v}
		} else {
			sv.left.Add(v)
		}
	} else if cmp > 0 {
		if sv.right == nil {
			sv.right = &SetValue{value: v}
		} else {
			sv.right.Add(v)
		}
	}
}

// Len returns how many elements are in the branches
func (sv *SetValue) Len() int {
	ret := 1
	if sv.left != nil {
		ret += sv.left.Len()
	}
	if sv.right != nil {
		ret += sv.right.Len()
	}
	return ret
}

// Has checks if this branch contains the value v
func (sv *SetValue) Has(v string) bool {
	ret, _ := sv.find(v)
	return ret != nil
}

// String gets a string value of this branch
func (sv *SetValue) String() string {
	var ret string
	if sv.left != nil {
		ret += sv.left.String() + ", "
	}
	ret += "\"" + sv.value + "\""
	if sv.right != nil {
		ret += ", " + sv.right.String()
	}
	return ret
}

// Slice returns a string slice of all values in the branch
func (sv *SetValue) Slice() []string {
	var ret []string
	if sv.left != nil {
		ret = sv.left.Slice()
	}
	ret = append(ret, sv.value)
	if sv.right != nil {
		ret = append(ret, sv.right.Slice()...)
	}
	return ret
}

// Equal returns whether the given sets are the same.
func Equal(s1, s2 Set) bool {
	return s1.String() == s2.String()
}

// Subset returns whether s1 is a subset of s2.
func Subset(s1, s2 Set) bool {
	if s1.Len() == 0 || s2.Len() == 0 {
		return false
	}
	s1Sl := s1.Slice()
	for i := range s1Sl {
		if !s2.Has(s1Sl[i]) {
			return false
		}
	}
	return true
}

// Disjoint returns whether two sets _do not_ intersect
func Disjoint(s1, s2 Set) bool {
	s1Sl := s1.Slice()
	for i := range s1Sl {
		if s2.Has(s1Sl[i]) {
			return false
		}
	}
	return false
}

// Intersection finds elements that exist in both sets and makes a new
// set of them
func Intersection(s1, s2 Set) Set {
	var vals []string
	s1Sl := s1.Slice()
	for i := range s1Sl {
		if s2.Has(s1Sl[i]) {
			vals = append(vals, s1Sl[i])
		}
	}
	return NewFromSlice(vals)
}

// Union gets all elements in both sets and makes a new set with them.
func Union(s1, s2 Set) Set {
	var vals []string
	vals = append(vals, s1.Slice()...)
	vals = append(vals, s2.Slice()...)
	return NewFromSlice(vals)
}

// Difference returns a Set of all elements in s1 that aren't in s2
func Difference(s1, s2 Set) Set {
	var vals []string
	s1Sl := s1.Slice()
	for i := range s1Sl {
		if !s2.Has(s1Sl[i]) {
			vals = append(vals, s1Sl[i])
		}
	}
	return NewFromSlice(vals)
}

// SymmetricDifference returns all elements from s1 & s2 that occur in only one of the
// sets.
func SymmetricDifference(s1, s2 Set) Set {
	return Union(Difference(s1, s2), Difference(s2, s1))
}