tests/weighting_n_cutting_the_tree/cutting.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 и сортируем индекс
		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)
		})

		// Ищем ветку подходящую под лимит
		var branchID string
		for _, v := range index {
			if v.weight < threshold {
				branchID = v.id
				break
			}
		}
		if branchID == "" {
			return nil, errors.New("[GetNodePath]: threshold too small")
		}

		// Определяем путь к найденной ветке
		nodeID := tree[branchID].node.ID
		var path string
		if _, ok := tree[nodeID]; !ok {
			return nil, errors.New("[GetNodePath] node not exists")
		}
		if tree[nodeID].parent == nil {
			path = "/"
		} else {
			nodeID = tree[nodeID].parent.ID
			path = tree[nodeID].node.Name
			for path != "/" {
				if tree[nodeID].parent.Name == "/" {
					path = "/" + path
					break
				}
				nodeID = tree[nodeID].parent.ID
				path = tree[nodeID].node.Name + "/" + path
			}
		}

		// Вычитаем ветку из веса предков
		currentID := branchID
		if tree[currentID].parent != nil {
			currentID = tree[currentID].parent.ID
			for {
				tree[currentID].weight -= tree[branchID].weight
				if tree[currentID].parent == nil {
					break
				}
				currentID = tree[currentID].parent.ID
			}
		}

		// Вырезаем её из среза вместе со всеми потомками
		// Удаляем ссылки на ноду из родительской ноды
		if tree[branchID].parent != nil {
			parent := tree[branchID].parent
			m := slices.IndexFunc(parent.Children, func(node *Node) bool {
				return node.ID == branchID
			})
			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[branchID].node
		// Удаляем ветви вместе со всеми потомками
		stack := []string{branchID}
		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)
			}
		}

		// Добавляем найденную ветку к res
		res = append(res, SubTree{
			path: path,
			tree: node,
		})
	}
	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
		}
	}
	if branchID == "" {
		return "", 0
	}
	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 path != "/" {
		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
	}
}

func GetNode(tree map[string]*WeightedNode, index []TreeIndex, threshold uint16) (string, string, error) {
	// Для работы необходим актуальный индекс мапы, отсортированный по убыванию
	var branchID string
	for _, v := range index {
		if v.weight < threshold {
			branchID = v.id
			break
		}
	}
	if branchID == "" {
		return "", "", errors.New("threshold too small")
	}
	currentID := tree[branchID].node.ID
	nodeWeight := tree[branchID].weight

	// GetNodePath возвращает полный путь ноды от корня
	if _, ok := tree[currentID]; !ok {
		return "", "", errors.New("node not exists")
	}
	if tree[currentID].parent == nil {
		return tree[currentID].node.ID, "/", nil
	}
	currentID = tree[currentID].parent.ID
	path := tree[currentID].node.Name
	for {
		if tree[currentID].parent.Name == "/" {
			path = "/" + path
			break
		}
		currentID = tree[currentID].parent.ID
		path = tree[currentID].node.Name + "/" + path
	}
	// RecalculateParents вычитает вес вырезанной ветки из всех родительских веток
	if tree[currentID].parent == nil {
		return "", "", errors.New("it's the root node")
	}
	currentID = tree[currentID].parent.ID
	for {
		tree[currentID].weight -= nodeWeight
		if tree[currentID].parent == nil {
			break
		}
		currentID = tree[currentID].parent.ID
	}
	return tree[branchID].node.ID, path, nil

}

// 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
}