数据结构

背景

树结构是一种常见数据结构，树结构其实是一种特殊的链结构，现在学习一下最基本的生成和打印功能。

代码

创建及打印树结构，使用切片创建树，支持使用nil来剪枝，即部分分支不创建。

package main

import (
	"bytes"
	"fmt"
)

func main() {
	tree := new(Tree)
	tree.Create([]interface{}{1, 2, 3, nil, 5, 6, 7})
	fmt.Println(tree)
}

type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

type Tree struct {
	node *TreeNode
}

func (t *Tree) String() string {
	if t.node == nil {
		return ""
	}
	return string(t.print(t.node, 0, true))
}

// 打印树
func (t *Tree) print(node *TreeNode, level int, tail bool) []byte {
	var b bytes.Buffer
	if node == nil {
	} else {
		b.Write(t.print(node.Left, level+1, false))

		b.Write(bytes.Repeat([]byte("    "), level))
		if tail {
			b.Write([]byte("└── "))
		} else {
			b.Write([]byte("┌── "))
		}

		b.Write([]byte(fmt.Sprint(node.Val)))
		b.Write([]byte("
"))

		b.Write(t.print(node.Right, level+1, true))
	}
	return b.Bytes()
}
// 创建树
func (t *Tree) Create(list []interface{}) {
	// 没有数据则退出
	if len(list) == 0 {
		return
	}

	// 第一条数据为空则返回空节点
	if list[0] == nil {
		return
	}

	// 创建一个新节点
	rootNode := &TreeNode{Val: list[0].(int)}
	list = list[1:] // 移除第一个数据

	// 如果根节点不存在则赋值给根节点
	if t.node == nil {
		t.node = rootNode
	}

	// 使用队列
	var queue []*TreeNode
	queue = append(queue, t.node) // 根节点添加到队列中

	for len(queue) != 0 {
		cur := queue[0]
		queue = queue[1:]

		if len(list) == 0 {
			break
		}

		if list[0] != nil {
			cur.Left = &TreeNode{Val: list[0].(int)}
			queue = append(queue, cur.Left)
		}
		list = list[1:] // 移除第一个数据

		if len(list) == 0 {
			break
		}

		if list[0] != nil {
			cur.Right = &TreeNode{Val: list[0].(int)}
			queue = append(queue, cur.Right)
		}
		list = list[1:] // 移除第一个数据
	}
}