Binary Tree: Depth First Traversal

16th March 2021 at 4:48pm
Binary Tree

二叉树的深度优先遍历。有三种顺序:

  • 中序(inorder):root、左、右
    • 或 root、右、左,没有本质区别
  • 前序(preorder):左,root,右
  • 后序(postorder):左,右,root

这里先阐述算法思路,后面给出完整代码示例。算法用递归写很容易,这里讲述 不用递归、用栈实现 的方法。还有一些奇技淫巧,连栈都不需要,我没有深入了解。

中序

算法:

  1. Create an empty stack S.
  2. Initialize current node as root
  3. Push the current node to S and set current = current->left until current is NULL
  4. If current is NULL and stack is not empty then
    1. Pop the top item from stack.
    2. Print the popped item, set current = popped_item->right
    3. Go to step 3.
  5. If current is NULL and stack is empty then we are done.

重点在于如何定义一次循环的逻辑边界。在这里是这样定义的:

  1. 每次循环,给定一个头节点,将其一系列的左节点入栈
    1. 如果这个头节点为空,跳过此步骤
  2. 从栈 Pop 出一个节点,把其添加到遍历结果中
  3. 把该节点的右节点列为下次循环的头节点

重点在于加粗部分的逻辑,把连续 pop 给表达了进来。每次都把某一批左左左节点入栈。

前序

有两种做法:

  1. 可以跟中序类似,区别仅在 push 元素进结果列表(answer)的时机不同;见下面代码中的 preorderTraversal()
  2. 更简洁的做法,从 root 入栈开始,每个循环 pop 出一个元素,并先后将其右节点、左节点入栈;见下面代码中的 preorderTraversal2()

后序

后序是最难的部分。因为后序的遍历过程不是个尾递归问题(未理解)。这篇 文章 描述 用两个栈 的思路,非常好理解。它的思路是,把遍历过程反过来,把后面要访问的节点先入栈,最终再从栈里出一个个 pop 出来。

看看代码实现就能理解。

代码

下面的代码以这颗树为测试用例:

       1
      / \
     2   3
    / \   \
   4   5   6
  /       / \
 7       8   9
              \
              10
package main

import "fmt"

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

func inorderTraversal(root *TreeNode) []int {
    answer := []int{}
    stack := []*TreeNode{}

    curr := root
    for curr != nil || len(stack) != 0 {
        for curr != nil {
            stack = append(stack, curr)
            curr = curr.Left
        }
        curr, stack = stack[len(stack)-1], stack[:len(stack)-1]
        answer = append(answer, curr.Val)
        curr = curr.Right
    }

    return answer
}

// Solution 1: Inorder 版稍加修改
func preorderTraversal(root *TreeNode) []int {
    answer := []int{}
    stack := []*TreeNode{}

    curr := root
    for curr != nil || len(stack) != 0 {
        for curr != nil {
            answer = append(answer, curr.Val)
            stack = append(stack, curr)
            curr = curr.Left
        }
        curr, stack = stack[len(stack)-1], stack[:len(stack)-1]
        curr = curr.Right
    }

    return answer
}

// Solution 2: 更好理解的方式
func preorderTraversal2(root *TreeNode) []int {
    answer := []int{}
    if root == nil {
        return answer
    }

    stack := []*TreeNode{root}
    for len(stack) != 0 {
        curr := stack[len(stack)-1]
        stack = stack[:len(stack)-1]
        answer = append(answer, curr.Val)

        if curr.Right != nil {
            stack = append(stack, curr.Right)
        }
        if curr.Left != nil {
            stack = append(stack, curr.Left)
        }
    }

    return answer
}

func postorderTraversal(root *TreeNode) []int {
    answer := []int{}
    firstStack := []*TreeNode{root}
    secondStack := []*TreeNode{}

    if root == nil {
        return answer
    }

    for len(firstStack) != 0 {
        curr := firstStack[len(firstStack)-1]
        firstStack = firstStack[:len(firstStack)-1]
        secondStack = append(secondStack, curr)

        if curr.Left != nil {
            firstStack = append(firstStack, curr.Left)
        }
        if curr.Right != nil {
            firstStack = append(firstStack, curr.Right)
        }
    }

    for len(secondStack) != 0 {
        answer = append(answer, secondStack[len(secondStack)-1].Val)
        secondStack = secondStack[:len(secondStack)-1]
    }
    return answer
}

func main() {
    root := &TreeNode{1,
        &TreeNode{2,
            &TreeNode{4, &TreeNode{7, nil, nil}, nil},
            &TreeNode{5, nil, nil}},
        &TreeNode{3, nil,
            &TreeNode{6,
                &TreeNode{8, nil, nil},
                &TreeNode{9,
                    nil,
                    &TreeNode{10, nil, nil},
                },
            }}}
    fmt.Println(inorderTraversal(root))
    fmt.Println(preorderTraversal(root))
    fmt.Println(preorderTraversal2(root))
    fmt.Println(postorderTraversal(root))
}

参考