业务代码很多都需要生成树结构数据返回给前端，正好Go的泛型还没咋学习。
所以就写了个泛型方案，以后那些需要返回树结构的业务都可以共用同一个泛型生成方案了。
```go
package main
import (
"encoding/json"
"fmt"
"sort"
)
func main() {
err := test()
if err != nil {
panic(err)
}
}
func test() error {
var (
// 一般从数据库查出来的都是切片数据,这里模拟数据库的数据
td = []*MyTree{
{1, 0, "top"},
{2, 1, "1"},
{21, 2, "11"},
{3, 1, "2"},
{31, 3, "21"},
}
// 指定排序规则
less = func(a, b *MyTree) bool { return a.Id < b.Id }
// 根据数据创建树结构
mt = MakeTree(td, less)
)
// 大多数情况,都是序列化为json返回给前端
b, err := json.MarshalIndent(mt, "", " ")
if err != nil {
return err
}
fmt.Printf("%T\n%s\n", mt, b)
st := RestoreTree(mt, less)
for i, v := range st {
fmt.Printf("index:%d,val:%#v\n", i, v)
}
return nil
}
// MyTree 这里是具体业务结构
type MyTree struct {
Id int `json:"id"`
ParentId int `json:"parentId"`
Data string `json:"data"`
}
func (t *MyTree) IsTop() bool {
return t.ParentId == 0
}
func (t *MyTree) GetId() int {
return t.Id
}
func (t *MyTree) GetParentId() int {
return t.ParentId
}
// TreeData 下面是泛型的通用代码
// 所有结构体都可以按照下面方式创建树结构,泛型结构体必须有下面几个方法
type TreeData interface {
IsTop() bool
GetId() int
GetParentId() int
}
type Tree[T TreeData] struct {
Data T `json:"data"` // 树结构数据内容
Children []*Tree[T] `json:"children"` // 树结构子节点,null表示没有子节点
}
func MakeTree[T TreeData](data []T, less func(a, b T) bool) (res *Tree[T]) {
treeMap := make(map[int]*Tree[T], len(data))
for _, v := range data {
treeMap[v.GetId()] = &Tree[T]{Data: v}
}
for _, val := range treeMap {
if val.Data.IsTop() {
res = val // 找到树的顶级节点
} else if p, ok := treeMap[val.Data.GetParentId()]; ok {
p.Children = append(p.Children, val)
if less != nil {
for i := len(p.Children) - 2; i >= 0; i-- {
if less(p.Children[i].Data, p.Children[i+1].Data) {
break // 已经有序
}
p.Children[i], p.Children[i+1] = p.Children[i+1], p.Children[i]
}
}
}
}
return
}
func RestoreTree[T TreeData](tree *Tree[T], less func(a, b T) bool) (res []T) {
var restore func(*Tree[T])
restore = func(tree *Tree[T]) {
if tree != nil {
res = append(res, tree.Data)
for _, v := range tree.Children {
restore(v)
}
}
}
restore(tree)
// 递归还原,并按照传入规则排序
if less != nil {
sort.Slice(res, func(i, j int) bool { return less(res[i], res[j]) })
}
return
}
```
有疑问加站长微信联系（非本文作者）
