刷题笔记：二叉树剪枝（递归，迭代）

题目

剑指 Offer II 047. 二叉树剪枝

首解（递归）

遇到这种问题果然第一反应还是递归，算是比较简单暴力的解法。
当左子树为空且右子树为空且本身值为0时，删除该节点（删除的方式就是返回nil）

/** * Definition for a binary tree node. * type TreeNode struct { * Val int * Left *TreeNode * Right *TreeNode * } */
func pruneTree(root *TreeNode) *TreeNode {
    
    if root == nil{
    
        return nil
    }
    //左右根
    root.Left = pruneTree(root.Left)
    root.Right = pruneTree(root.Right)
    if root.Left == nil && root.Right == nil && root.Val == 0{
    
        return nil
    }
    return root
}

居然官方也没给出什么其他解法。那么就用显式的栈来实现一下吧。

再解：构造栈

注意在遍历完某一节点的左子树之后，我们会遍历其右子树，此时要标记这个节点我们已经访问完左子树了，否则会一直循环访问其右子树。

/** * Definition for a binary tree node. * type TreeNode struct { * Val int * Left *TreeNode * Right *TreeNode * } */
func pruneTree(node *TreeNode) *TreeNode {
    
    // 造一个栈
    stack := []*TreeNode{
    } // 这里保存指向节点的指针而不是本身
    visited := new(TreeNode) // 防止无限访问右子树
    root := node
    for root != nil || len(stack)> 0{
    
        // 先遍历左
        for root != nil{
    
            // 根入栈
            stack = append(stack, root)
            root = root.Left
        }
        // fmt.Println(root.Val)
        // 此时root为空
        // 左遍历完了，从栈中取出根，遍历其右子树
        if len(stack) > 0{
    
            //取出根
            root = stack[len(stack) - 1]
            //尝试遍历右 跳回循环起始处
            if root.Right != nil && visited != root.Right{
     // 这里注意是root.right是否被访问过 如果访问过就不要再循环了。不要写成root
                root = root.Right
            } else {
    
                // 没有右，就出栈根
                stack = stack[:len(stack)-1]
                //注意这里出栈了，所以后面一定要判断栈是否为空 这里必须有判断root.right是否为空
                //否则会干掉子树有1的节点（上面可能因为visited而没有进入if 即已经访问过 但不一定为空
                if root.Val == 0 && root.Left == nil &&root.Right == nil{
    
                    // 将其父节点从栈中找出来，将指向root的指针置空
                    if len(stack) >0{
    
                        father := stack[len(stack)-1]
                        if father.Left == root{
    
                            father.Left = nil
                        } else if father.Right == root{
    
                            father.Right = nil
                        }
                    } else {
    
                        // 已经没有栈了，说明最后一个节点已经出栈，左右根遍历结束，且此节点应置为空
                        return nil
                    }
                }
                //左右根都遍历完了
                visited = root // 他不需要再遍历了，记录此节点
                root = nil
            }
        }
    }
    return node
}