golang tree 辅助操作

在项目开发过程中，发现需要很多树的数据结构，所以封装了一个树的接口，方便操作。

树的生成：

比如树的节点结构定义大概如下：

type struct Node {
	Name string
	Children []Node
	CalcRule string // 自定义属性
}

Name 指节点名字，Children是包含的子节点的数组，还有其他一些字段，就是叶子节点的一些属性了，比如这里的CalcRule

现在有一些节点的数据和他们对应的树的路径，比如：

A: {Name: "node1",  CalcRule: "add"}，对应的中间路径是"a" -> "b1" -> "c3";

B: {Name: "node2",  CalcRule: "minus"}，对应的中间路径是"a" -> "b1";

C: {Name: "node3",  CalcRule: "mult"}，对应的中间路径是"a" -> "b2" -> "c4";

D: {Name: "node4",  CalcRule: "add"}，对应的中间路径是"a" -> "b1" -> "c3";

E: {Name: "node5",  CalcRule: "none"}，对应的中间路径是"a" -> "b1" -> "c3";

按照这些路径生成好的树的结构图如下：

使用了接口创建这颗树的代码非常简单，如下所示：

func createTree() *Node {
	// 路径数组，每个路径需要包含叶子节点的名称，但是不需要包含根节点的名字
	pathArr := [][]string{
		[]string{b1", "c3", "node1"},
		[]string{b1", "node2"},
		[]string{b2", "c4", "node3"},
		[]string{b1", "c3", "node4"},
		[]string{b1", "c3", "node5"},
	}

	// 每个叶子节点的信息，在叶子节点产生之后赋值使用
	nodeArr := []Node{
		Node{CalcRule: "add"},
		Node{CalcRule: "minus"},
		Node{CalcRule: "mult"},
		Node{CalcRule: "add"},
		Node{CalcRule: "none"},
	}

	// 使用根节点信息初始化整棵树
	rootTree := &Node{
		Name: "a", // 根节点的名称放在这里
	}

	for k, v := range pathArr {
		// 通过路径信息和自定义的创建节点函数生成树的分支，返回的是叶子节点
		leaf := tree.CreateByPath(rootTree, pathArr, rootTree.CreateNode)

		// 增加叶子节点的信息
		leaf.(*Node).CalcRule = nodeArr[k].CalcRule
	}

	// 返回生成好的树
	return rootTree
}

通过tree.CreateByPath函数生成树，参数中rootTree.CreateNode是Node结构的成员函数。为了使用tree.CreateByPath函数，Node结构必须实现NormalTree（自己起的名字）的接口，还需要实现其他几个成员函数，如下所示：

// 满足接口的函数：获取到子节点数组，返回的数组元素必须要实现NormalTree的接口
// 因为children数组的元素是Node，而*Node才实现了NormalTree接口，所以需要将[]Node转换为[]*Node
func (self *Node) GetChildren() []interface{} {
	if nil == self.Children || 0 == len(self.Children) {
		return nil
	}

	arr := make([]interface{}, len(self.Children))
	for k := range self.Children {
		arr[k] = &self.Children[k]
	}

	return arr
}

// 满足接口的函数：增加子节点，其中使用了自定义的create函数。返回结果是满足NormalTree接口的子节点
func (self *Node) AddChild(node map[string]interface{}, create func(map[string]interface{}) interface{}) interface{} {
	child := create(node).(Node)
	self.Children = append(self.Children, child)

	return &self.Children[len(self.Children)-1]
}

// 满足接口的函数：获取节点名字
func (self *Node) GetName() string {
	return self.Name
}

// 在AddChild中使用的create函数，参数data为固定的map:{"value": 节点名称}。返回结果是新的节点
func (self *Node) CreateNode(data map[string]interface{}) interface{} {
	node := Node{
		Name:    data["value"].(string),
	}

	return node
}

下面是tree接口文件tree.go的代码：

package tree

type NormalTree interface {
	GetChildren() []interface{}
	AddChild(map[string]interface{}, func(map[string]interface{}) interface{}) interface{}
	GetName() string
}

func CreateByPath(tree NormalTree, path []string, creatFunc func(map[string]interface{}) interface{}) NormalTree {
	for _, v := range path {
		index := -1
		children := tree.GetChildren()
		for k1, v1 := range children {
			vData := v1.(NormalTree)
			if vData.GetName() == v {
				index = k1
				break
			}
		}

		if -1 == index {
			index = len(children)
			tree.AddChild(map[string]interface{}{"value": v}, creatFunc)
			children = tree.GetChildren()
		}

		tree = children[index].(NormalTree)
	}

	return tree
}

从路径数组生成树的工作就完成了。

有时候我也需要从遍历生成的树，每次写代码用广度优先搜素或者深度优先搜素也很麻烦。所以加了一个将树转换为路径链表的功能。

按照上面生成好的树，转换回来就如下所示：

[]Node{{node2, minus} -> {b1} ->{a},  {node1, add} -> {c3} -> {b1} -> {a}, ......} // 用伪代码，大家理解意思就行。

为了满足这个操作，Node结构还需要增加一个成员和实现另外一个成员函数：

type struct Node {
    Name string
    Children []Node
    Parent *Node // 增加指向父节点的指针
    CalcRule string // 自定义属性
}

func (self *Node) SetParent(parent interface{}) {
	self.Parent = parent.(*Node)
}

tree.go文件中接口定义也需要增加这个函数定义：

SetParent(interface{})

做这个操作的tree.go中函数如下：

func GetReversePath(tree NormalTree) []NormalTree {
	arr := []NormalTree{tree}
	result := make([]NormalTree, 0)

	for len(arr) > 0 {
		next := make([]NormalTree, 0)

		for k, v := range arr {
			children := v.GetChildren()
			if nil != children && len(children) > 0 {
				for k1, v1 := range children {
					v1Data := v1.(NormalTree)
					children[k1].(NormalTree).SetParent(arr[k])
					next = append(next, v1Data)
				}
			} else {
				result = append(result, v)
			}
		}

		arr = next
	}

	return result
}

最后，在加一个树的拷贝函数，可以将一个实现了NormalTree接口的结构A的树拷贝为路径相同的实现了NormalTree的结构B的树。因为这个使用的比较少，就不细说了，把tree.go中的函数列下即可：

type NodeContent interface {
	Decode() map[string]interface{}
}

unc Copy(src NormalTree, dst NormalTree, creatFunc func(map[string]interface{}) interface{}, level int) {
	srcArr := []NormalTree{src}
	dstArr := []NormalTree{dst}

	currentLevel := 1
	for len(srcArr) > 0 && (level <= 0 || currentLevel < level) {
		nextSrcArr := make([]NormalTree, 0)
		nextDstArr := make([]NormalTree, 0)

		for k, v := range srcArr {
			srcChildren := v.GetChildren()
			if nil == srcChildren || 0 == len(srcChildren) {
				continue
			}

			for _, v1 := range srcChildren {
				dstChild := dstArr[k].AddChild(v1.(NodeContent).Decode(), creatFunc)
				if nil == dstChild {
					continue
				}

				v1Data := v1.(NormalTree)
				dstData := dstChild.(NormalTree)

				nextSrcArr = append(nextSrcArr, v1Data)
				nextDstArr = append(nextDstArr, dstData)
			}

		}

		srcArr = nextSrcArr
		dstArr = nextDstArr

		currentLevel += 1
	}
}