the-cat-tree-fun.go

the-cat-tree-fun.go

package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"strings"
)

const (
	k_InvalidArgumentException = "InvalidArgumentException"
)

func main() {

	var categoryTree = CategoryTree{
		Set: []*Node{},
	}
	categoryTree.Set = append(categoryTree.Set, &Node{
		Name:     "A",
		Children: []*Node{},
	})
	categoryTree.Set = append(categoryTree.Set, &Node{
		Name:     "B",
		Children: []*Node{},
	})

	// DUMP
	fmt.Printf("SET: %s", categoryTree)

	categoryTree.Add("A", "C")
	categoryTree.Add("A", "D")
	categoryTree.Add("A", "E")

	categoryTree.Add("B", "F")
	categoryTree.Add("B", "G")
	categoryTree.Add("B", "H")

	categoryTree.Add("C", "I")

	categoryTree.Add("H", "J")

	// DUMP
	fmt.Printf("SET: %s", categoryTree)
}

type CategoryTree struct {
	Set []*Node `json:"set"`
}

func (thisRef *CategoryTree) Add(parentName string, childName string) error {
	// Check if anything exists
	for _, rootNode := range thisRef.Set {
		if rootNode.Search(childName) != nil {
			return errors.New(k_InvalidArgumentException)
		}
	}

	// Find the parent to add to
	var parentToAddTo *Node
	for _, rootNode := range thisRef.Set {
		var theNodeToAddTo = rootNode.Search(parentName)
		if theNodeToAddTo != nil {
			parentToAddTo = theNodeToAddTo
			break
		}
	}

	if parentToAddTo == nil {
		return errors.New(k_InvalidArgumentException)
	}

	parentToAddTo.Children = append(parentToAddTo.Children, &Node{
		Name:     childName,
		Children: []*Node{},
	})

	return nil
}

// String - Stringer interface
func (thisRef CategoryTree) String() string {
	data, err := json.Marshal(thisRef)
	if err != nil {
		log.Fatal(err)
	}

	return string(data)
}

type Node struct {
	Name     string  `json:"name"`
	Children []*Node `json:"children"`
}

// Search - Searches for a node by name
func (thisRef *Node) Search(name string) *Node {
	// ~ Execute DFS - Pre-order

	if strings.Compare(thisRef.Name, name) == 0 {
		return thisRef
	}

	for _, childNode := range thisRef.Children {
		node := childNode.Search(name)
		if node != nil {
			return node
		}
	}

	return nil
}