Thomas Ngo

Software Engineer

3 minute read

Problem Statement

problem-1325

Intuition

To solve the problem of removing leaf nodes with a specific target value from a binary tree, we need to recursively check each node’s children. If a node becomes a leaf node and its value equals the target value, we need to remove it. This removal process should continue until no such nodes remain.

Approach

  • Recursive Traversal: Traverse the tree using post-order traversal (left, right, root) because we need to handle the children before their parent. This ensures that we can safely remove leaf nodes and check if their parent becomes a leaf node after their removal.
  • Remove Leaf Nodes: If a node is a leaf node and its value is equal to the target value, we can remove it by returning None from the recursive call.
  • Return Updated Tree: The final tree is returned after all necessary removals.

Complexity

  • Time complexity: O(n) where n is the number of nodes in the tree

  • Space complexity: O(h) where h is the height of the tree

Code

# 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 removeLeafNodes(self, root: Optional[TreeNode], target: int) -> Optional[TreeNode]:
        def helper(node):
            if not node:
                return

            L = helper(node.left)
            R = helper(node.right)
            if node.val == target:
                if not L and not R:
                    node.val = 'x'
                if not L and R and R.val == 'x':
                    node.val = 'x'
                if L and L.val == 'x' and not R:
                    node.val = 'x'
                if L and R and L.val == 'x' and R.val == 'x':
                    node.val = 'x'
            return node

        def delete_node(node):
            if not node:
                return

            if node.left and node.left.val == 'x':
                node.left = None
            if node.right and node.right.val == 'x':
                node.right = None

            delete_node(node.left)
            delete_node(node.right)


        helper(root)
        delete_node(root)

        return root if root.val != 'x' else None

Editorial Solution

Approach 1: Recursion (Postorder Traversal)

class Solution:
    def removeLeafNodes(self, root: Optional[TreeNode], target: int) -> Optional[TreeNode]:
        # Base case
        if root is None:
            return None

        # 1. Visit the left children
        root.left = self.removeLeafNodes(root.left, target)

        # 2. Visit the right children
        root.right = self.removeLeafNodes(root.right, target)

        # 3. Check if the current node is a leaf node and its value equals target
        if root.left is None and root.right is None and root.val == target:
            # Delete the node by returning None to the parent, effectively disconnecting it
            return None

        # Keep it untouched otherwise
        return root

Approach 2: Iterative (PostOrder Traversal)

class Solution:
    def removeLeafNodes(self, root: Optional[TreeNode], target: int) -> Optional[TreeNode]:
        if not root:
            return None

        stack = []
        current_node = root
        last_right_node = None

        while stack or current_node:
            # Push left nodes to the stack until reaching the leftmost node.
            while current_node:
                stack.append(current_node)
                current_node = current_node.left

            # Access the top node on the stack without removing it.
            # This node is the current candidate for processing.
            current_node = stack[-1]

            # Check if the current node has an unexplored right subtree.
            # If so, shift to the right subtree unless it's the subtree we just visited.
            if current_node.right and current_node.right is not last_right_node:
                current_node = current_node.right
                continue  # Continue in the while loop to push this new subtree's leftmost nodes.

            # Remove the node from the stack, since we're about to process it.
            stack.pop()

            # If the node has no right subtree or the right subtree has already been visited,
            # then check if it is a leaf node with the target value.
            if not current_node.right and not current_node.left and current_node.val == target:
                # If the stack is empty after popping, it means the root was a target leaf node.
                if not stack:
                    return None  # The tree becomes empty as the root itself is removed.

                # Identify the parent of the current node.
                parent = stack[-1] if stack else None

                # Disconnect the current node from its parent.
                if parent and parent.left is current_node:
                    parent.left = None  # If current is a left child, set the left pointer to null.
                elif parent and parent.right is current_node:
                    parent.right = None  # If current is a right child, set the right pointer to null.

            # Mark this node as visited by setting 'last_right_node' to 'current_node' before moving to the next iteration.
            last_right_node = current_node
            # Reset 'current_node' to None to ensure the next node from the stack is processed.
            current_node = None

        return root  # Return the modified tree

You may also enjoy