feat: add tree

tree/tree.go

@@ -0,0 +1,66 @@
+package tree
+
+import (
+	"cmp"
+)
+
+type NodeSlice[I cmp.Ordered, E any] []NodeTree[I, E]
+
+func (nodes NodeSlice[T, E]) Len() int {
+	return len(nodes)
+}
+func (nodes NodeSlice[T, E]) Less(i, j int) bool {
+	return nodes[i].GetId() < nodes[j].GetId()
+}
+func (nodes NodeSlice[T, E]) Swap(i, j int) {
+	nodes[i], nodes[j] = nodes[j], nodes[i]
+}
+
+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
+}
+
+// IntoTree 列表转树, 切片无children, map转换
+// 两次循环就可以获取列表转树
+func IntoTree[I cmp.Ordered, E Node[I, U], U NodeTree[I, U]](rows []E, rootPid I) []U {
+	nodes := make([]U, 0, len(rows))
+	nodeMaps := make(map[I]U)
+	for _, v := range rows {
+		vv := v.MapTree()
+		nodes = append(nodes, vv)
+		nodeMaps[v.MapId()] = vv
+	}
+	return intoTree(nodeMaps, nodes, rootPid)
+}
+
+// IntoTree 列表转树
+// 两次循环就可以获取列表转树
+func IntoTree2[T cmp.Ordered, E NodeTree[T, E]](rows []E, rootPid T) []E {
+	nodeMaps := make(map[T]E)
+	for _, v := range rows {
+		vv := v
+		nodeMaps[v.GetId()] = vv
+	}
+	return intoTree(nodeMaps, rows, rootPid)
+}
+
+func intoTree[T cmp.Ordered, E NodeTree[T, E]](nodeMaps map[T]E, rows []E, rootPid T) []E {
+	var root []E
+	for _, v := range rows {
+		node := v
+		pid := node.GetPid()
+		if pid == rootPid {
+			root = append(root, node)
+		} else if parent, exists := nodeMaps[pid]; exists {
+			parent.AppendChildren(node)
+		}
+	}
+	return root
+}

tree/tree_test.go

@@ -0,0 +1,55 @@
+package tree
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type Dept struct {
+	Id   int
+	Pid  int
+	Name string
+}
+
+func (d *Dept) MapId() int { return d.Id }
+func (d *Dept) MapTree() *DeptTree {
+	return &DeptTree{
+		Dept:     d,
+		Children: nil,
+	}
+}
+
+type DeptTree struct {
+	*Dept
+	Children []*DeptTree
+}
+
+func (d *DeptTree) GetId() int  { return d.Id }
+func (d *DeptTree) GetPid() int { return d.Pid }
+
+func (d *DeptTree) AppendChildren(v *DeptTree) {
+	d.Children = append(d.Children, v)
+}
+
+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"},
+}
+
+func TestTree(t *testing.T) {
+	vv := IntoTree(arr, 0)
+	v, err := json.MarshalIndent(vv, " ", "  ")
+	require.NoError(t, err)
+	t.Log(string(v))
+}