45 lines
1011 B
Go
45 lines
1011 B
Go
package main
|
|
|
|
import (
|
|
"github.com/google/uuid"
|
|
"maps"
|
|
)
|
|
|
|
type WeightedTree struct {
|
|
Tree *Node
|
|
Weights map[string]*WeightedNode
|
|
}
|
|
|
|
func getWeightedTree(branches, depth int, parent *Node) WeightedTree {
|
|
if depth == 0 {
|
|
return WeightedTree{}
|
|
}
|
|
id := uuid.NewString()
|
|
node := &Node{
|
|
ID: id,
|
|
Name: id,
|
|
Children: []*Node{},
|
|
}
|
|
tree := WeightedTree{
|
|
Tree: node,
|
|
Weights: map[string]*WeightedNode{id: &WeightedNode{
|
|
node: node,
|
|
parent: parent,
|
|
children: nil,
|
|
weight: node.getNodeWeight(),
|
|
}},
|
|
}
|
|
for i := 0; i < branches; i++ {
|
|
childTree := getWeightedTree(branches, depth-1, node)
|
|
if childTree.Tree != nil {
|
|
tree.Tree.Children = append(tree.Tree.Children, childTree.Tree)
|
|
tree.Weights[id].children = append(tree.Weights[id].children, childTree.Tree)
|
|
maps.Copy(tree.Weights, childTree.Weights)
|
|
}
|
|
}
|
|
for _, child := range tree.Tree.Children {
|
|
tree.Weights[id].weight = tree.Weights[id].weight + tree.Weights[child.ID].weight
|
|
}
|
|
return tree
|
|
}
|