二叉树递归改迭代

参考二叉树八股文：递归改迭代 :: labuladong的算法小抄 (gitee.io)

递归二叉树很简单

func traverse(root *TreeNode) {
    if root == nil {
        return nil
    }
    traverse(root.Left)
    // 中序
    traverse(root.Right)
}

用栈来模拟递归

// 引入栈
var stack = list.New()
func traverse(root *TreeNode) {
    if root == nil {
        return
    }
    stack.PushBack(root)
    traverse(root.Left)
    traverse(root.Right)
    stack.Remove(stack.Back())
}

// 去除递归
// 基本思路是
// 对根节点，一路把它和左节点全部压栈
// 重复：弹出节点，检查右节点是否非空，非空就继续把左节点压栈
var stack = list.New()

func traverse(root *TreeNode) {
    pushLeftBranch(root)
    for stack.Len() != 0 {
        p := stack.Remove(stack.Back()).(*TreeNode)
        pushLeftBranch(p.Right)
    }
}

func pushLeftBranch(p *TreeNode) {
    for p != nil {
        stack.PushBack(p)
        p = p.Left
    }
}

但是这样子依旧不能判断前中后三种遍历的位置。

核心在哪呢？

在于拿到一个节点p的时候，判断它的左右节点有没有遍历过。比如说中序，就是当前节点p的左节点已经遍历，而右节点还没有遍历。

var stack = list.New()

func traverse(root *TreeNode) {
    visited := &TreeNode{}
    pushLeftBranch(root)
    for stack.Len() != 0 {
        p := stack.Back().Value.(*TreeNode)

        if (p.Left == nil || p.Left == visited) && p.Right != visited {
            // 中序遍历位置：保证右节点没访问过，且左节点已访问过(左节点为空或visited)
            pushLeftBranch(p.Right)
        }

        if p.Right == nil || p.Right == visited {
            // 后序遍历位置: 保证右节点已经访问。左节点的已访问由上面的中序那里来完成。
            visited = stack.Remove(stack.Back()).(*TreeNode)
        }
    }
}

func pushLeftBranch(p *TreeNode) {
    for p != nil {
        // 前序遍历位置：即将被压入栈，意味着子节点都没被遍历
        stack.PushBack(p)
        p = p.Left
    }
}