Skip to content

Commit

Permalink
fix: fix tree
Browse files Browse the repository at this point in the history
  • Loading branch information
thinkgos committed Nov 22, 2024
1 parent e582dae commit 21ba114
Show file tree
Hide file tree
Showing 3 changed files with 51 additions and 27 deletions.
34 changes: 11 additions & 23 deletions tree/tree.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,17 +4,7 @@ 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
Expand All @@ -27,39 +17,37 @@ type Node[I cmp.Ordered, U NodeTree[I, U]] interface {
MapTree() U
}

// IntoTree 列表转树, 切片无children, map转换
// IntoTree 列表转树, 切片无children
// 两次循环就可以获取列表转树
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
e := v.MapTree()
nodes = append(nodes, e)
nodeMaps[v.MapId()] = e
}
return intoTree(nodeMaps, nodes, rootPid)
}

// IntoTree 列表转树
// IntoTree 列表转树, 切片有children
// 两次循环就可以获取列表转树
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
for _, e := range rows {
nodeMaps[e.GetId()] = e
}
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()
for _, e := range rows {
pid := e.GetPid()
if pid == rootPid {
root = append(root, node)
root = append(root, e)
} else if parent, exists := nodeMaps[pid]; exists {
parent.AppendChildren(node)
parent.AppendChildren(e)
}
}
return root
Expand Down
15 changes: 15 additions & 0 deletions tree/tree_sort.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,15 @@
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]
}
29 changes: 25 additions & 4 deletions tree/tree_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,13 +7,19 @@ import (
"github.com/stretchr/testify/require"
)

var _ Node[int, *DeptTree] = (*Dept)(nil)
var _ NodeTree[int, *DeptTree] = (*DeptTree)(nil)

type Dept struct {
Id int
Pid int
Name string
}

// MapId implements Node.
func (d *Dept) MapId() int { return d.Id }

// MapTree implements Node.
func (d *Dept) MapTree() *DeptTree {
return &DeptTree{
Dept: d,
Expand All @@ -26,9 +32,13 @@ type DeptTree struct {
Children []*DeptTree
}

func (d *DeptTree) GetId() int { return d.Id }
// GetId implements NodeTree.
func (d *DeptTree) GetId() int { return d.Id }

// GetPid implements NodeTree.
func (d *DeptTree) GetPid() int { return d.Pid }

// AppendChildren implements NodeTree.
func (d *DeptTree) AppendChildren(v *DeptTree) {
d.Children = append(d.Children, v)
}
Expand All @@ -48,8 +58,19 @@ var arr = []*Dept{
}

func TestTree(t *testing.T) {
vv := IntoTree(arr, 0)
v, err := json.MarshalIndent(vv, " ", " ")
gotTree1 := IntoTree(arr, 0)
tree1, err := json.MarshalIndent(gotTree1, " ", " ")
require.NoError(t, err)

arrTree := make([]*DeptTree, 0, len(arr))
for _, v := range arr {
vv := *v
arrTree = append(arrTree, vv.MapTree())
}

gotTree2 := IntoTree2(arrTree, 0)
tree2, err := json.MarshalIndent(gotTree2, " ", " ")
require.NoError(t, err)
t.Log(string(v))

require.Equal(t, tree2, tree1)
}

0 comments on commit 21ba114

Please sign in to comment.