i2-bill-erp/utils/build_tree.go

package utils

import (
	"container/list"
	"git.i2edu.net/i2/i2-bill-erp/transform"
)

func BuildTree(nodes *[]*transform.TreeNode) ([]*transform.TreeNode, error) {
	originNodeList := list.New()
	treeNodeList := list.New()
	var rootArr []*transform.TreeNode
	for _, node := range *nodes {
		if node.Parent == 0 {
			treeNodeList.PushBack(node)
			rootArr = append(rootArr, node)
		} else {
			originNodeList.PushBack(node)
		}
	}

	// 把子节点根据parent分配到对应的父节点上
	for ele := treeNodeList.Front(); ele != nil; ele = ele.Next() {
		treeNode := ele.Value.(*transform.TreeNode)
		originEle := originNodeList.Front()
		if originEle == nil {
			break
		}
		for originEle != nil {
			originNextEle := originEle.Next()
			originNode := originEle.Value.(*transform.TreeNode)
			if originNode.Parent == treeNode.Id {
				treeNodeList.InsertAfter(originNode, ele)
				treeNode.Nodes = append(treeNode.Nodes, originNode)

				originNodeList.Remove(originEle)
			}
			originEle = originNextEle
		}
	}
	return rootArr, nil
}