exercism/go/custom-set/stringset.go

474 lines
9.1 KiB
Go
Raw Normal View History

2016-08-15 19:08:39 +00:00
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 {
2016-08-16 18:58:10 +00:00
top *Node
2016-08-15 19:08:39 +00:00
}
2016-08-16 18:58:10 +00:00
type Node struct {
2016-08-15 19:08:39 +00:00
value string
2016-08-16 18:58:10 +00:00
left *Node
right *Node
2016-08-15 19:08:39 +00:00
}
// New returns an empty Set
func New() Set {
s := new(Set)
2016-08-16 18:58:10 +00:00
s.top = &Node{}
2016-08-15 19:08:39 +00:00
return *s
}
// NewFromSlice takes a slice of strings and returns a Set
func NewFromSlice(s []string) Set {
ret := New()
2016-08-16 18:58:10 +00:00
r := &ret
2016-08-15 19:08:39 +00:00
for i := range s {
2016-08-16 18:58:10 +00:00
r.Add(s[i])
2016-08-15 19:08:39 +00:00
}
2016-08-16 18:58:10 +00:00
return *r
2016-08-15 19:08:39 +00:00
}
// Add adds a value to the set
func (s Set) Add(v string) {
if s.top.value == "" {
s.top.value = v
return
}
2016-08-16 18:58:10 +00:00
if !s.Has(v) {
newNode := s.top.Add(&Node{value: v})
2016-08-23 20:47:19 +00:00
2016-08-16 18:58:10 +00:00
s.top.value = newNode.value
s.top.left = newNode.left
s.top.right = newNode.right
2016-08-15 19:08:39 +00:00
}
}
// 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
2016-08-16 18:58:10 +00:00
var delNode, repNode, orphan *Node
2016-08-15 19:08:39 +00:00
if sv == nil {
// Deleting 'top'
2016-08-16 18:58:10 +00:00
delNode = s.top
2016-08-15 19:08:39 +00:00
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
}
2016-08-16 18:58:10 +00:00
s.top = repNode
2016-08-15 19:08:39 +00:00
} else {
cmp = strings.Compare(v, sv.value)
if cmp < 0 && sv.left != nil {
// It's the left node
delNode = sv.left
} else if cmp > 0 && sv.right != nil {
// 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)
}
}
}
2016-08-16 18:58:10 +00:00
func (s Set) PrettyPrint() {
fmt.Println("=========")
s.pp(s.top, 0)
fmt.Println("=========")
}
func (s Set) pp(n *Node, indent int) {
if n.left == n || n.right == n {
// Circular reference...
fmt.Println("Circular reference... You done messed up.")
return
}
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)
}
}
func (s *Set) find(v string) *Node {
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) (*Node, 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 *Node) {
if s.top.value == "" {
s.top.value = v.value
s.top.left = v.left
s.top.right = v.right
return
}
s.top.findHome(v)
}
2016-08-15 19:08:39 +00:00
// 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.
2016-08-16 18:58:10 +00:00
func (sv *Node) find(val string) (*Node, error) {
2016-08-15 19:08:39 +00:00
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.
2016-08-16 18:58:10 +00:00
func (sv *Node) findParent(val string) (*Node, error) {
2016-08-15 19:08:39 +00:00
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")
}
2016-08-16 18:58:10 +00:00
func (sv *Node) findHome(v *Node) {
2016-08-15 19:08:39 +00:00
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)
}
}
}
2016-08-16 18:58:10 +00:00
func (sv *Node) Add(n *Node) *Node {
cmp := strings.Compare(n.value, sv.value)
2016-08-15 19:08:39 +00:00
if cmp < 0 {
if sv.left == nil {
2016-08-16 18:58:10 +00:00
sv.left = n
2016-08-15 19:08:39 +00:00
} else {
2016-08-16 18:58:10 +00:00
sv.left = sv.left.Add(n)
2016-08-15 19:08:39 +00:00
}
2016-08-16 18:58:10 +00:00
} else {
2016-08-15 19:08:39 +00:00
if sv.right == nil {
2016-08-16 18:58:10 +00:00
sv.right = n
2016-08-15 19:08:39 +00:00
} else {
2016-08-16 18:58:10 +00:00
sv.right = sv.right.Add(n)
2016-08-15 19:08:39 +00:00
}
}
2016-08-16 18:58:10 +00:00
2016-08-23 20:47:19 +00:00
newVal := sv
/* TODO: Fix the tree balancing
2016-08-16 18:58:10 +00:00
balance := GetBalance(sv)
// left left unbalance
2016-08-23 20:47:19 +00:00
if balance > 1 && strings.Compare(n.value, sv.getLeftValue()) < 0 {
2016-08-16 18:58:10 +00:00
newVal = RightRotate(sv)
}
// right right unbalance
2016-08-23 20:47:19 +00:00
if balance < -1 && strings.Compare(n.value, sv.getRightValue()) > 0 {
2016-08-16 18:58:10 +00:00
newVal = LeftRotate(sv)
}
// left right unbalance
2016-08-23 20:47:19 +00:00
if balance > 1 && strings.Compare(n.value, sv.getLeftValue()) > 0 {
2016-08-16 18:58:10 +00:00
sv.left = LeftRotate(sv.left)
newVal = RightRotate(sv)
}
// right left unbalance
2016-08-23 20:47:19 +00:00
if balance < -1 && strings.Compare(n.value, sv.getRightValue()) < 0 {
2016-08-16 18:58:10 +00:00
sv.right = RightRotate(sv.right)
newVal = LeftRotate(sv)
}
2016-08-23 20:47:19 +00:00
*/
2016-08-16 18:58:10 +00:00
return newVal
}
func (sv *Node) getLeftValue() string {
if sv.left == nil {
return ""
}
return sv.left.value
}
func (sv *Node) getRightValue() string {
if sv.right == nil {
return ""
}
return sv.right.value
2016-08-15 19:08:39 +00:00
}
// Len returns how many elements are in the branches
2016-08-16 18:58:10 +00:00
func (sv *Node) Len() int {
if sv == nil {
return 0
}
2016-08-15 19:08:39 +00:00
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
2016-08-16 18:58:10 +00:00
func (sv *Node) Has(v string) bool {
2016-08-15 19:08:39 +00:00
ret, _ := sv.find(v)
return ret != nil
}
// String gets a string value of this branch
2016-08-16 18:58:10 +00:00
func (sv *Node) String() string {
2016-08-15 19:08:39 +00:00
var ret string
if sv.left != nil {
2016-08-16 18:58:10 +00:00
if sv.left == sv {
ret += "**, "
} else {
ret += sv.left.String() + ", "
}
2016-08-15 19:08:39 +00:00
}
ret += "\"" + sv.value + "\""
if sv.right != nil {
2016-08-16 18:58:10 +00:00
if sv.right == sv {
ret += "**, "
} else {
ret += ", " + sv.right.String()
}
2016-08-15 19:08:39 +00:00
}
return ret
}
// Slice returns a string slice of all values in the branch
2016-08-16 18:58:10 +00:00
func (sv *Node) Slice() []string {
2016-08-15 19:08:39 +00:00
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 {
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
}
}
2016-08-16 18:58:10 +00:00
return true
2016-08-15 19:08:39 +00:00
}
// 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))
}
2016-08-16 18:58:10 +00:00
/* Helper Functions for Balancing the Tree */
func RightRotate(y *Node) *Node {
x := y.left
t2 := x.right
// Perform rotation
x.right = y
y.left = t2
// Return new root
return x
}
func LeftRotate(x *Node) *Node {
y := x.right
t2 := y.left
// Perform rotation
y.left = x
x.right = t2
// Return new root
return y
}
func GetBalance(n *Node) int {
if n == nil {
return 0
}
return n.left.Len() - n.right.Len()
}
func CompareNodes(n1, n2 *Node) int {
return strings.Compare(n1.value, n2.value)
}