使用go语言将slice(list)数据转换为tree数据

241 阅读1分钟

之前在迭代中有一个需求,要求接口返回的数据是tree型数据,故而在这记录下相关的代码及后续的改进

下面给出一个基于go语言的解决方案
type Node struct {
	id       string
	pid      string
	children []Node
}

func main() {
        // 初始化要挂接的数据
	nodes := []Node{{id: "4", pid: "3"}, {id: "3", pid: "2"}, {id: "2", pid: "1"}}
	// 要展示的tree数据
        tree := make([]Node, 0)
        // 根节点,叶子节点
	var roots, children []Node
        // “1”对应的是根节点
	for _, v := range nodes {
		if v.pid == "1" {
			roots = append(roots, v)
		} else {
			children = append(children, v)
		}
	}

	for _, v := range roots {
		rootNode := &Node{
			id:  v.id,
			pid: v.pid,
		}
		recursion(rootNode, children)
		tree = append(tree, *rootNode)
	}
	fmt.Print(tree)
}
func recursion(rootNode *Node, children []Node) {
	id := rootNode.id
	for _, v := range children {
		if id == v.pid {
			childNode := &Node{
				id:  v.id,
				pid: v.pid,
			}
			recursion(childNode, children)
			rootNode.children = append(rootNode.children, *childNode)
		}
	}
}

可以看出这段代码的children即使已经挂接在父节点下也是每次都在循环。那么有没有办法每次只去挂接父节点对应的子节点呢?有的!我这里用了map,给数据加一个索引,方便挂接,话不多说给出代码
type Node struct {
	id       string
	pid      string
	children []Node
}

func main() {
	nodes := []Node{{id: "4", pid: "3"}, {id: "3", pid: "2"}, {id: "2", pid: "1"}}
	tree := make([]Node, 0)

	rootMap := make(map[string]Node)
	childrenMap := make(map[string][]Node)

	for _, v := range nodes {
		if v.pid == "1" {
			rootMap[v.id] = v
		} else {
			childrenMap[v.pid] = append(childrenMap[v.pid], v)
		}
	}

	for _, v := range rootMap {
		rootNode := &Node{
			id:  v.id,
			pid: v.pid,
		}
		recursion(rootNode, childrenMap)
		tree = append(tree, *rootNode)
	}
	fmt.Print(tree)
}
func recursion(rootNode *Node, childrenMap map[string][]Node) {
	treeValue, ok := childrenMap[rootNode.id]
	if ok {
		for _, v := range treeValue {
			childNode := &Node{
				id:  v.id,
				pid: v.pid,
			}
			recursion(childNode, childrenMap)
			rootNode.children = append(rootNode.children, *childNode)
		}
	}
}

有不对或者更好的方法希望大家可以提出来!!!