Leetcode 二叉树生成及转换成数组的工具包

我们 Leetcode 刷题应该是每个程序员小伙伴都会经历的一个状态，但当每次遇到二叉树相关的题目时，实现函数虽然写完了，想自己在本地 IDE 中做好充分的验证的话，还是会有一些麻烦的。

尤其是写测试用例时是根据 Leetcode 官方的数组格式进行设置传入二叉树结构，但是数组类型转成二叉树结构就写的有点蛋疼了。也就是二叉树结构的序列化为数组和数组反序列化为二叉树结构问题了，总不能每次遇到时临时写一坨坨的代码吧，所以就有了下面基于 Golang 语言实现的一个简单版本。

为啥要在本地写单元测试，Leetcode 有控制台，贴上去运行呀。的确可以这么做，但还是不方便，并且你也不想每次提交之后都是未通过吧，这样会显得我们写的代码很烂[机制boy]。

话不多说，上代码先：

二叉树标准结构

// TreeNode Definition for a binary tree node
type TreeNode struct {
    Val   int       `json:"val"`
    Left  *TreeNode `json:"left,omitempty"`
    Right *TreeNode `json:"right,omitempty"`
}

自定义的 Null 整形常量

// Null 自定义常量`Null[MinInt32]`标识的元素解释为nil节点
const Null int = math.MinInt32

计算二叉树最大深度

// Depth 获取二叉树最大深度
func Depth(root *TreeNode) int {
    var depthF func(root *TreeNode) float64
    depthF = func(root *TreeNode) float64 {
        if root == nil {
            return 0
        }

        return math.Max(depthF(root.Left), depthF(root.Right)) + 1
    }

    return int(depthF(root))
}

二叉树数组反序列化为二叉树

// Unmarshal 将整形切片反序列化为标准二叉树结构对象，其中自定义常量`Null[MinInt32]`标识的元素解释为nil节点
//
// 输入：nums = [5,4,8,11,null,13,4,7,2,null,null,5,1]
//
// 输出（图形化结构）:
//
//                5
//              /   \
//             4     8
//            / \   / \
//           11 nl 13  4
//          / \   / \  / \
//         7   2 nl nl 5  1
//
func Unmarshal(nums []int) *TreeNode {
    if len(nums) == 0 {
        return nil
    }

    root := new(TreeNode)
    root.Val = nums[0]

    queue := list.New()
    queue.PushFront(root)

    for i := 0; queue.Len() > 0 && i < len(nums); {
        element := queue.Back()
        queue.Remove(element)

        cur := element.Value.(*TreeNode)
        if cur.Val == Null {
            continue
        }

        if i++; i < len(nums) {
            left := new(TreeNode)
            left.Val = nums[i]
            if left.Val != Null {
                cur.Left = left
            }
            queue.PushFront(left)
        }

        if i++; i < len(nums) {
            right := new(TreeNode)
            right.Val = nums[i]
            if right.Val != Null {
                cur.Right = right
            }
            queue.PushFront(right)
        }
    }

    return root
}

二叉树序列化为二叉树数组

// Marshal 将标准二叉树结构对象序列化为整形切片，其中nil节点会按自定义常量`Null[MinInt32]`标识
func Marshal(root *TreeNode) (res []int) {
    depth := Depth(root)

    var dfs func(depth int, nodes ...*TreeNode)
    dfs = func(curDepth int, nodes ...*TreeNode) {
        if curDepth == 1 && len(nodes) == 1 && nodes[0] == nil {
            return
        }
        if len(nodes) <= 0 {
            return
        }

        nexts := []*TreeNode{}
        for _, node := range nodes {
            if node == nil {
                res = append(res, Null)
                continue
            }

            res = append(res, node.Val)
            if curDepth < depth {
                nexts = append(nexts, node.Left)
                nexts = append(nexts, node.Right)
            }
        }
        dfs(curDepth+1, nexts...)
    }

    dfs(1, root)
    return
}

二叉树数组格式化输出Leetcode字符串结构

// Format2Leetcode 根据 Leetcode 测试用例的输入/输出数据格式进行格式化
func Format2Leetcode(nums []int) string {
    buffer := bytes.NewBuffer(nil)
    buffer.WriteByte('[')
    for i, num := range nums {
        if num == math.MinInt32 {
            buffer.WriteString("null")
        } else {
            buffer.WriteString(fmt.Sprint(num))
        }
        if i < len(nums)-1 {
            buffer.WriteByte(',')
        }
    }
    buffer.WriteByte(']')

    return string(buffer.Bytes())
}

测试 Example

以 Leetcode - 226 题为例进行测试。

func Example() {
    // Leetcode
    // 226. 翻转二叉树
    // 期望输入 [4,2,7,1,3,6,9]
    // 期望输出 [4,7,2,9,6,3,1]

    var invertTree func(root *TreeNode) *TreeNode
    invertTree = func(root *TreeNode) *TreeNode {
        if root == nil {
            return root
        }
        root.Left, root.Right = root.Right, root.Left
        invertTree(root.Left)
        invertTree(root.Right)

        return root
    }

    // 定义输入切片
    nums := []int{4, 2, 7, 1, 3, 6, 9}
    // 实际输入
    actualInput := Format2Leetcode(nums)
    fmt.Println("实际输入：", actualInput)

    // 反序列化为二叉树结构
    root := Unmarshal(nums)
    // 翻转二叉树算法函数实现
    invertRoot := invertTree(root)
    // 序列化二叉树为切片
    invertNums := Marshal(invertRoot)

    // 实际输出
    actualOutput := Format2Leetcode(invertNums)
    fmt.Println("实际输出：", actualOutput)

    // Output:
    //
    // 实际输入： [4,2,7,1,3,6,9]
    // 实际输出： [4,7,2,9,6,3,1]
}

完整代码地址

该二叉树生成及转换成数组的工具包已经上传到 Github 的 leetcode-toolkit。