LC 0814 [M] Binary Tree Pruning - ALawliet/algorithms GitHub Wiki
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, val=0, left=None, right=None):
# self.val = val
# self.left = left
# self.right = right
class Solution:
def pruneTree(self, root: Optional[TreeNode]) -> Optional[TreeNode]:
def postorder(node):
if not node:
return 0
L = postorder(node.left)
R = postorder(node.right)
# print(f'{L}/{node.val}\{R}')
if L == 0:
node.left = None
if R == 0:
node.right = None
return L + R + node.val
postorder(root)
'''
edge case
Input: [0,null,0,0,0]
Output: [0]
Expected: []
'''
if not root.left and not root.right and root.val == 0:
return None
return root
def pruneTree(self, root):
if not root: return None
root.left = self.pruneTree(root.left)
root.right = self.pruneTree(root.right)
if not root.left and not root.right and not root.val: return None
return root
class Solution:
def pruneTree(self, root):
def dfs(node):
if not node: return True
left, right = dfs(node.left), dfs(node.right)
if left: node.left = None
if right: node.right = None
return left and right and node.val == 0
return root if not dfs(root) else None