Leetcode 1796 Solution

This article provides solution to leetcode question 1796 (correct-a-binary-tree)

https://leetcode.com/problems/correct-a-binary-tree

Solution

# 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 correctBinaryTree(self, root: TreeNode) -> TreeNode:
        curr_nodes = {root: None}

        while curr_nodes:
            found_invalid = False

            for node, parent in curr_nodes.items():
                if node.left not in curr_nodes and node.right not in curr_nodes:
                    continue

                if parent.left == node:
                    parent.left = None
                    return root
                elif parent.right == node:
                    parent.right = None
                    return root
                else:
                    raise Exception("It can't be possible!")

            next_nodes = {}
            for node in curr_nodes.keys():
                if node.left:
                    next_nodes[node.left] = node
                if node.right:
                    next_nodes[node.right] = node

            curr_nodes = next_nodes

        return root