tree-building

tree_building.go

package tree

import (
	"errors"
	"sort"
)

type Record struct {
	ID int
	Parent int
}

type Node struct {
	ID int
	Children []*Node
}

var nodes = make(map[int]*Node)

func Build(records []Record) (*Node, error) {
	if len(records) < 1 {
		return nil, nil
	}
	var root *Node
	sort.SliceStable(records, func(i, j int) bool { return records[i].ID < records[j].ID })
	if records[0].ID != 0 {
		return nil, errors.New("no parent")
	}
	var r Record
	r, records = records[0], records[1:]
	if r.Parent != 0 {
		return nil, errors.New("root has parent")
	}
	root = &Node{ID: 0}
	nodes[0] = root
	var prev Record
	for _, r := range records {
		if prev.ID + 1 != r.ID {
			return nil, errors.New("non-continuous")
		}

		if r.Parent == 0 && r.ID == 0 {
			return nil, errors.New("duplicate root")
		}

		if containsDuplicate(records, r) {
			return nil, errors.New("duplicate node")
		}

		if r.ID == r.Parent || r.ID < r.Parent {
			return nil, errors.New("cycle")
		}

		if parent, ok := nodes[r.Parent]; !ok {
			child := root.insertChild(r.Parent)
			child.insertChild(r.ID)
			nodes[r.Parent] = child
		} else {
			child := parent.insertChild(r.ID)
			nodes[r.ID] = child
		}
		prev = r
	}
	return root, nil
}

func containsDuplicate(records []Record, record Record) bool {
	count := 0
	for _, r := range records {
		if r.ID == record.ID {
			count++
		}
	}
	return count > 1
}

func (n *Node) insertChild(ID int) *Node {
	child := Node{
		ID:       ID,
		Children: nil,
	}
	n.Children = append(n.Children, &child)
	return &child
}