[二叉树] 二叉树剪枝

前言

剪枝操作是也算二叉树的一个基本操作之一，包括回溯算法等，剪枝的思想都是算法优化的一个重要考量，今天记录一下这道题。

题目

二叉树剪枝

给定二叉树根结点 root ，此外树的每个结点的值要么是 0，要么是 1。

返回移除了所有不包含 1 的子树的原二叉树。

( 节点 X 的子树为 X 本身，以及所有 X 的后代。)

题目拆解

这道题是一个剪枝的基本题目，没有太多复杂的逻辑，所以重点就放在剪枝的dfs上就好。

1、一趟深度优先
2、剪枝的判定是当前节点的所有子树中，有没有1，没有就剪掉，有就留着。

简单来说，这两步就够了，还有几个细节可能需要注意，当前节点的判断条件要根据子节点返回来确定，所以用后续遍历的结构来处理。

代码

class Solution {
    public TreeNode pruneTree(TreeNode root) {
        boolean dfs = dfs(root);
        if (!dfs) return null;
        return root;
    }

    boolean dfs(TreeNode node) {
        // 如果当前节点为null 返回false
        if (node == null) return false;
        // 先往下找
        // 如果不包含1,直接删掉
        boolean leftRes = dfs(node.left);
        if (!leftRes) node.left = null;
        boolean rightRes = dfs(node.right);
        if (!rightRes) node.right = null;

        // 如果左右子树中有1,当前节点不删除
        if(leftRes || rightRes) return true;
        // 如果左右子树都没有1,那么根据当前节点值是否等于1 决定返回结果.
        return node.val == 1;
    }
}

总结

这题也是少有的一次bugfree就通过的leecode题目，通过代码也可以看出来，逻辑并不复杂，可以整理一下当做剪枝的标准模板使用了。