1660

1660. Correct a Binary Tree (Medium)

You have a binary tree with a small defect. There is exactly one invalid node where its right child incorrectly points to another node at the same depth but to the invalid node's right.

Given the root of the binary tree with this defect, root, return the root of the binary tree after removing this invalid node and every node underneath it (minus the node it incorrectly points to).

Custom testing:

The test input is read as 3 lines:

• TreeNode root
• int fromNode (not available to correctBinaryTree)
• int toNode (not available to correctBinaryTree)

After the binary tree rooted at root is parsed, the TreeNode with value of fromNode will have its right child pointer pointing to the TreeNode with a value of toNode. Then, root is passed to correctBinaryTree.

 

Example 1:

Input: root = [1,2,3], fromNode = 2, toNode = 3
Output: [1,null,3]
Explanation: The node with value 2 is invalid, so remove it.

Example 2:

Input: root = [8,3,1,7,null,9,4,2,null,null,null,5,6], fromNode = 7, toNode = 4
Output: [8,3,1,null,null,9,4,null,null,5,6]
Explanation: The node with value 7 is invalid, so remove it and the node underneath it, node 2.

 

Constraints:

• The number of nodes in the tree is in the range [3, 104].
• -109 <= Node.val <= 109
• All Node.val are unique.
• fromNode != toNode
• fromNode and toNode will exist in the tree and will be on the same depth.
• toNode is to the right of fromNode.
• fromNode.right is null in the initial tree from the test data.

Companies: Google

Related Topics:
Hash Table, Tree, Depth-First Search, Breadth-First Search, Binary Tree

Similar Questions:

• Flatten Binary Tree to Linked List (Medium)
• Flatten a Multilevel Doubly Linked List (Medium)

Hints:

• If you traverse the tree from right to left, the invalid node will point to a node that has already been visited.

Solution 1. BFS

// OJ: https://leetcode.com/problems/correct-a-binary-tree/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
class Solution {
public:
    TreeNode* correctBinaryTree(TreeNode* root) {
        queue<TreeNode*> q{{root}};
        unordered_map<TreeNode*, TreeNode*> parent;
        while (q.size()) {
            int cnt = q.size();
            while (cnt--) {
                auto n = q.front();
                q.pop();
                if (parent.count(n->right)) {
                    auto p = parent[n];
                    if (n == p->left) p->left = nullptr;
                    else p->right = nullptr;
                } else {
                    if (n->left) {
                        q.push(n->left);
                        parent[n->left] = n;
                    }
                    if (n->right) {
                        q.push(n->right);
                        parent[n->right] = n;
                    }
                }
            }
        }
        return root;
    }
};

Solution 2. DFS from Right to Left

// OJ: https://leetcode.com/problems/correct-a-binary-tree/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(N)
class Solution {
    unordered_set<TreeNode*> seen;
public:
    TreeNode* correctBinaryTree(TreeNode* root) {
        if (!root || seen.count(root->right)) return nullptr;
        seen.insert(root);
        root->right = correctBinaryTree(root->right);
        root->left = correctBinaryTree(root->left);
        return root;
    }
};