fix: add tree sort

tree/sort.go

@@ -0,0 +1 @@
+package tree

tree/tree.go

@@ -2,45 +2,85 @@ package tree
 
 import (
 	"cmp"
+	"slices"
 )
 
+type Node[I cmp.Ordered, U NodeTree[I, U]] interface {
+	MapId() I
+	MapTree() U
+}
+
 type NodeTree[I cmp.Ordered, T any] interface {
 	GetId() I
 	GetPid() I
 	AppendChildren(T)
 }
 
-type Node[I cmp.Ordered, U NodeTree[I, U]] interface {
-	MapId() I
-	MapTree() U
+type NodeSort[T any] interface {
+	SortChildren(cmp func(a, b T) int)
 }
 
-// IntoTree 列表转树, 切片无children
-// 两次循环就可以获取列表转树
-func IntoTree[I cmp.Ordered, E Node[I, U], U NodeTree[I, U]](rows []E, rootPid I) []U {
+// Map implement Node map to NodeTree
+func Map[I cmp.Ordered, E Node[I, U], U NodeTree[I, U]](rows []E) []U {
 	nodes := make([]U, 0, len(rows))
-	nodeMaps := make(map[I]U)
 	for _, v := range rows {
+		nodes = append(nodes, v.MapTree())
+	}
+	return nodes
+}
+
+// IntoTree 列表转树, 切片无children
+// 元素顺序由x本身顺序决定, 可提前排序, 然后转树(或使用 SortFunc)
+func IntoTree[T cmp.Ordered, E Node[T, U], U NodeTree[T, U]](x []E, rootPid T) []U {
+	nodeMaps, nodes := intoMapTree(x)
+	return intoTree(nodeMaps, nodes, rootPid)
+}
+
+// IntoTree 列表转树, 切片有children, 顺序由x本身顺序决定
+// 元素顺序由x本身顺序决定, 可提前排序, 然后转树(或使用 SortFunc)
+func IntoTree2[T cmp.Ordered, E NodeTree[T, E]](x []E, rootPid T) []E {
+	nodeMaps := intoMap(x)
+	return intoTree(nodeMaps, x, rootPid)
+}
+
+// SortFunc 树排序
+func SortFunc[T NodeSort[T]](x []T, cmp func(a, b T) int) {
+	if len(x) == 0 {
+		return
+	}
+	slices.SortFunc(x, cmp)
+	for _, v := range x {
+		v.SortChildren(cmp)
+	}
+}
+
+// T -> U 的映射
+// E -> U 转换
+func intoMapTree[T cmp.Ordered, E Node[T, U], U NodeTree[T, U]](x []E) (map[T]U, []U) {
+	nodes := make([]U, 0, len(x))
+	nodeMaps := make(map[T]U)
+	for _, v := range x {
 		e := v.MapTree()
 		nodes = append(nodes, e)
 		nodeMaps[v.MapId()] = e
 	}
-	return intoTree(nodeMaps, nodes, rootPid)
+	return nodeMaps, nodes
 }
 
-// IntoTree 列表转树, 切片有children
-// 两次循环就可以获取列表转树
-func IntoTree2[T cmp.Ordered, E NodeTree[T, E]](rows []E, rootPid T) []E {
+// T -> E 映射
+func intoMap[T cmp.Ordered, E NodeTree[T, E]](x []E) map[T]E {
+	// T -> E 映射
 	nodeMaps := make(map[T]E)
-	for _, e := range rows {
+	for _, e := range x {
 		nodeMaps[e.GetId()] = e
 	}
-	return intoTree(nodeMaps, rows, rootPid)
+	return nodeMaps
 }
 
-func intoTree[T cmp.Ordered, E NodeTree[T, E]](nodeMaps map[T]E, rows []E, rootPid T) []E {
+// 转树
+func intoTree[T cmp.Ordered, E NodeTree[T, E]](nodeMaps map[T]E, x []E, rootPid T) []E {
 	var root []E
-	for _, e := range rows {
+	for _, e := range x {
 		pid := e.GetPid()
 		if pid == rootPid {
 			root = append(root, e)
@@ -50,12 +90,3 @@ func intoTree[T cmp.Ordered, E NodeTree[T, E]](nodeMaps map[T]E, rows []E, rootP
 	}
 	return root
 }
-
-// Map implement Node map to NodeTree
-func Map[I cmp.Ordered, E Node[I, U], U NodeTree[I, U]](rows []E) []U {
-	nodes := make([]U, 0, len(rows))
-	for _, v := range rows {
-		nodes = append(nodes, v.MapTree())
-	}
-	return nodes
-}

tree/tree_test.go

@@ -9,11 +9,13 @@ import (
 
 var _ Node[int, *DeptTree] = (*Dept)(nil)
 var _ NodeTree[int, *DeptTree] = (*DeptTree)(nil)
+var _ NodeSort[*DeptTree] = (*DeptTree)(nil)
 
 type Dept struct {
 	Id   int
 	Pid  int
 	Name string
+	Sort int
 }
 
 // MapId implements Node.
@@ -43,18 +45,27 @@ func (d *DeptTree) AppendChildren(v *DeptTree) {
 	d.Children = append(d.Children, v)
 }
 
+// SortChildren implements NodeSort.
+func (d *DeptTree) SortChildren(cmp func(a, b *DeptTree) int) {
+	SortFunc(d.Children, cmp)
+}
+
+func CmpDept(a, b *DeptTree) int {
+	return a.Sort - b.Sort
+}
+
 var arr = []*Dept{
-	{1, 0, "超然科技"},
-	{8, 3, "bb"},
-	{6, 2, "研究院"},
-	{2, 0, "低速科技"},
-	{3, 1, "科研中心"},
-	{4, 1, "运营中心"},
-	{5, 2, "吃喝院"},
-	{7, 3, "aa"},
-	{9, 4, "cc"},
-	{10, 5, "dd"},
-	{11, 6, "ee"},
+	{1, 0, "超然科技", 1},
+	{8, 3, "bb", 2},
+	{6, 2, "研究院", 4},
+	{2, 0, "低速科技", 5},
+	{3, 1, "科研中心", 4},
+	{4, 1, "运营中心", 2},
+	{5, 2, "吃喝院", 3},
+	{7, 3, "aa", 4},
+	{9, 4, "cc", 2},
+	{10, 5, "dd", 3},
+	{11, 6, "ee", 4},
 }
 
 func TestTree(t *testing.T) {
@@ -68,4 +79,9 @@ func TestTree(t *testing.T) {
 	require.NoError(t, err)
 
 	require.Equal(t, tree2, tree1)
+
+	t.Log(string(tree1))
+	SortFunc(gotTree1, CmpDept)
+	tree11, _ := json.MarshalIndent(gotTree1, " ", "  ")
+	t.Log(string(tree11))
 }