泛型方法创建树结构数据

业务代码很多都需要生成树结构数据返回给前端，正好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 } ```