Lstmxx的空间

814. 二叉树剪枝

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

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

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

链接

/**
 * Definition for a binary tree node.
 * class TreeNode {
 *     val: number
 *     left: TreeNode | null
 *     right: TreeNode | null
 *     constructor(val?: number, left?: TreeNode | null, right?: TreeNode | null) {
 *         this.val = (val===undefined ? 0 : val)
 *         this.left = (left===undefined ? null : left)
 *         this.right = (right===undefined ? null : right)
 *     }
 * }
 */

function pruneTree(root: TreeNode | null): TreeNode | null {
  function dfs(node: TreeNode | null): boolean {
    if (node === null) return false
    const l = dfs(node.left)
    const r = dfs(node.right)
    if (!l) node.left = null
    if (!r) node.right = null
    return node.val === 1 || l || r
  }
  const res = dfs(root)
  return res ? root : null
};