File size: 2,339 Bytes
6380833 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | package treex
type Tree[T any] struct {
root *Node[T]
}
// NewTree attempts to create a new Tree from a root node.
// It traverses the children graph to ensure it is acyclic.
func NewTree[T any](root *Node[T]) (*Tree[T], error) {
if root == nil {
return nil, ErrRootNodeIsNil
}
visited := make(map[*Node[T]]bool) // visited nodes tracker (as we don't yet know if it's a tree or not)
onStack := make(map[*Node[T]]bool) // to verify graph is acyclic
var dfs func(n *Node[T]) error
dfs = func(n *Node[T]) error {
if n == nil {
return ErrNodeGraphInvalid
}
if onStack[n] {
return ErrGraphHasCycle
}
if visited[n] {
return nil
}
visited[n] = true
onStack[n] = true
for _, child := range n.children {
if child == nil {
return ErrNodeGraphInvalid
}
if err := dfs(child); err != nil {
return err
}
}
onStack[n] = false
return nil
}
if err := dfs(root); err != nil {
return nil, err
}
return &Tree[T]{root: root}, nil
}
func (t *Tree[T]) Root() *Node[T] {
return t.root
}
func (t *Tree[T]) DFS(cb func(n *Node[T]) (stop bool, err error)) error {
if t.root == nil {
return nil
}
var walk func(n *Node[T]) error
walk = func(n *Node[T]) error {
stop, err := cb(n)
if err != nil {
return err
}
if stop {
return nil
}
for _, child := range n.children {
if child == nil {
return ErrNodeGraphInvalid
}
if err := walk(child); err != nil {
return err
}
}
return nil
}
return walk(t.root)
}
// Leafs returns all leaf nodes in the tree
func (t *Tree[T]) Leafs() []*Node[T] {
leafs := make([]*Node[T], 0)
_ = t.DFS(func(n *Node[T]) (bool, error) {
if n.IsLeaf() {
leafs = append(leafs, n)
}
return false, nil
})
return leafs
}
// SwapNode can swap any node in the tree including the root.
// If the old node wasn't found an error is returned.
// The new node's parent will be set to the old node's parent.
//
// CAUTION: if you swap a node while walking the tree, you should start backtracking after the swap
// otherwise you'll keep iterating the detached subtree
func (t *Tree[T]) SwapNode(old *Node[T], new *Node[T]) error {
if old == t.Root() {
t.root = new
return nil
}
if old.Parent() == nil {
return ErrNodeHasNoParentButNotRoot
}
parent := old.Parent()
return parent.SwapChild(old, new)
}
|