137 lines
4.1 KiB
Go
137 lines
4.1 KiB
Go
package main
|
||
|
||
import (
|
||
"github.com/pkg/errors"
|
||
"slices"
|
||
)
|
||
|
||
func DecomposeTree(tree map[string]*WeightedNode, threshold uint16) ([]SubTree, error) {
|
||
var res []SubTree
|
||
// Рабочий цикл
|
||
for len(tree) > 0 {
|
||
// Индексируем tree по weight и сортируем индекс
|
||
index := GetTreeIndex(tree)
|
||
// Ищем ветку подходящую под лимит
|
||
branchID, weight := GetBranch(tree, index, threshold)
|
||
|
||
// Определяем путь к найденной ветке
|
||
path, err := GetNodePath(tree, branchID)
|
||
if err != nil {
|
||
return nil, errors.Wrap(err, "[GetNodePath] ")
|
||
}
|
||
// Вычитаем ветку из веса предков
|
||
RecalculateParents(tree, branchID, weight)
|
||
|
||
// Вырезаем её из среза вместе со всеми потомками
|
||
branch := CutBranch(tree, branchID)
|
||
// Добавляем найденную ветку к res
|
||
res = append(res, SubTree{
|
||
path: path,
|
||
tree: branch,
|
||
})
|
||
}
|
||
return res, nil
|
||
}
|
||
|
||
type TreeIndex struct {
|
||
id string
|
||
weight uint16
|
||
}
|
||
|
||
// GetTreeIndex индексирует дерево по weight и сортирует индекс по уменьшению weight
|
||
func GetTreeIndex(tree map[string]*WeightedNode) []TreeIndex {
|
||
var index = make([]TreeIndex, len(tree))
|
||
i := 0
|
||
for k, v := range tree {
|
||
index[i] = TreeIndex{
|
||
id: k,
|
||
weight: v.weight,
|
||
}
|
||
i += 1
|
||
}
|
||
slices.SortFunc(index, func(a, b TreeIndex) int {
|
||
return int(b.weight) - int(a.weight)
|
||
})
|
||
return index
|
||
}
|
||
|
||
// GetBranch возвращает первую найденную ветвь дерева, суммарный вес которой меньше threshold
|
||
// Для работы необходим актуальный индекс мапы, отсортированный по убыванию
|
||
func GetBranch(tree map[string]*WeightedNode, index []TreeIndex, threshold uint16) (string, uint16) {
|
||
var branchID string
|
||
for _, v := range index {
|
||
if v.weight < threshold {
|
||
branchID = v.id
|
||
break
|
||
}
|
||
}
|
||
return tree[branchID].node.ID, tree[branchID].weight
|
||
}
|
||
|
||
// GetNodePath возвращает полный путь ноды от корня
|
||
func GetNodePath(tree map[string]*WeightedNode, nodeID string) (string, error) {
|
||
if _, ok := tree[nodeID]; !ok {
|
||
return "", errors.New("node not exists")
|
||
}
|
||
if tree[nodeID].parent == nil {
|
||
return "/", nil
|
||
}
|
||
nodeID = tree[nodeID].parent.ID
|
||
path := tree[nodeID].node.Name
|
||
for {
|
||
if tree[nodeID].parent.Name == "/" {
|
||
path = "/" + path
|
||
break
|
||
}
|
||
nodeID = tree[nodeID].parent.ID
|
||
path = tree[nodeID].node.Name + "/" + path
|
||
}
|
||
return path, nil
|
||
}
|
||
|
||
// RecalculateParents вычитает вес вырезанной ветки из всех родительских веток
|
||
func RecalculateParents(tree map[string]*WeightedNode, currentID string, nodeWeight uint16) {
|
||
if tree[currentID].parent == nil {
|
||
return
|
||
}
|
||
currentID = tree[currentID].parent.ID
|
||
for {
|
||
tree[currentID].weight -= nodeWeight
|
||
if tree[currentID].parent == nil {
|
||
break
|
||
}
|
||
currentID = tree[currentID].parent.ID
|
||
}
|
||
}
|
||
|
||
// CutBranch удаляет из дерева ветку со всеми потомками, а так же удаляет ссылку на себя из родительской ноды
|
||
func CutBranch(tree map[string]*WeightedNode, nodeID string) *Node {
|
||
// Удаление ссылки на ноду из родительской ноды
|
||
if tree[nodeID].parent != nil {
|
||
parent := tree[nodeID].parent
|
||
m := slices.IndexFunc(parent.Children, func(node *Node) bool {
|
||
return node.ID == nodeID
|
||
})
|
||
if m >= 0 {
|
||
parent.Children[m] = parent.Children[len(parent.Children)-1]
|
||
parent.Children = parent.Children[:len(parent.Children)-1]
|
||
//parent.Children = append(parent.Children[:m], parent.Children[m+1:]...)
|
||
}
|
||
}
|
||
|
||
node := tree[nodeID].node
|
||
// Удаление ветви вместе со всеми потомками
|
||
stack := []string{nodeID}
|
||
for len(stack) > 0 {
|
||
current := stack[len(stack)-1]
|
||
stack = stack[:len(stack)-1]
|
||
if _, ok := tree[current]; ok {
|
||
for _, v := range tree[current].children {
|
||
stack = append(stack, v.ID)
|
||
}
|
||
delete(tree, current)
|
||
}
|
||
}
|
||
return node
|
||
}
|