814. 二叉树剪枝 : 简单递归运用题

题目描述

这是 LeetCode 上的 814. 二叉树剪枝 ，难度为中等。

Tag : 「二叉树」、「DFS」、「递归」

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

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

节点 node 的子树为 node 本身加上所有 node 的后代。

示例 1：

输入：root = [1,null,0,0,1]

输出：[1,null,0,null,1]

解释：
只有红色节点满足条件“所有不包含 1 的子树”。 右图为返回的答案。

示例 2：

输入：root = [1,0,1,0,0,0,1]

输出：[1,null,1,null,1]

示例 3：

输入：root = [1,1,0,1,1,0,1,0]

输出：[1,1,0,1,1,null,1]

提示：

树中节点的数目在范围 $[1, 200]$ 内
Node.val 为 $0$ 或 $1$

递归

根据题意，我们将原函数 pruneTree 作为递归函数，递归函数的含义为「将入参 root 中的所有不包含 $1$ 的子树移除，并返回新树头结点」。

不失一般性的考虑任意节点作为入参该如何处理：我们可以递归处理左右子树，并将新左右子树重新赋值给 root。由于当前节点 root 的左右子树可能为空树，因此我们要增加递归函数入参为空的边界处理。

当递归操作完成后，若左右节点任一值不为空（说明当前节点 root 不为叶子节点），我们可以直接返回 root，否则根据 root 的值是否为 $0$ 来决定返回空树还是 root 本身。

Java 代码：

class Solution {
    public TreeNode pruneTree(TreeNode root) {
        if (root == null) return null;
        root.left = pruneTree(root.left);
        root.right = pruneTree(root.right);
        if (root.left != null || root.right != null) return root;
        return root.val == 0 ? null : root;
    }
}

TypeScript 代码：

function pruneTree(root: TreeNode | null): TreeNode | null {
    if (root == null) return null
    root.left = pruneTree(root.left)
    root.right = pruneTree(root.right)
    if (root.left != null || root.right != null) return root
    return root.val == 0 ? null : root
};

时间复杂度： $O(n)$
空间复杂度：忽略递归带来的额外空间开销，复杂度为 $O(1)$

最后

这是我们「刷穿 LeetCode」系列文章的第 No.814 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：github.com/SharingSour… 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。

我正在参与掘金技术社区创作者签约计划招募活动，点击链接报名投稿。