Question

1660. Correct a Binary Tree

中文文档

Description

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.

Solutions

Solution 1: DFS

We design a function dfs(root) to handle the subtree with root as the root. If root is null or root.right has been visited, root is an invalid node, so we return null. Otherwise, we recursively process root.right and root.left, and return root.

Finally, we return dfs(root).

The time complexity is $O(n)$, and the space complexity is $O(n)$. Here, $n$ is the number of nodes in the binary tree.

# 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:
    def dfs(root):
      if root is None or root.right in vis:
        return None
      vis.add(root)
      root.right = dfs(root.right)
      root.left = dfs(root.left)
      return root

    vis = set()
    return dfs(root)

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *   int val;
 *   TreeNode left;
 *   TreeNode right;
 *   TreeNode() {}
 *   TreeNode(int val) { this.val = val; }
 *   TreeNode(int val, TreeNode left, TreeNode right) {
 *     this.val = val;
 *     this.left = left;
 *     this.right = right;
 *   }
 * }
 */
class Solution {
  private Set<TreeNode> vis = new HashSet<>();

  public TreeNode correctBinaryTree(TreeNode root) {
    return dfs(root);
  }

  private TreeNode dfs(TreeNode root) {
    if (root == null || vis.contains(root.right)) {
      return null;
    }
    vis.add(root);
    root.right = dfs(root.right);
    root.left = dfs(root.left);
    return root;
  }
}

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *   int val;
 *   TreeNode *left;
 *   TreeNode *right;
 *   TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *   TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *   TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
  TreeNode* correctBinaryTree(TreeNode* root) {
    unordered_set<TreeNode*> vis;
    function<TreeNode*(TreeNode*)> dfs = [&](TreeNode* root) -> TreeNode* {
      if (!root || vis.count(root->right)) {
        return nullptr;
      }
      vis.insert(root);
      root->right = dfs(root->right);
      root->left = dfs(root->left);
      return root;
    };
    return dfs(root);
  }
};

/**
 * Definition for a binary tree node.
 * function TreeNode(val, left, right) {
 *   this.val = (val===undefined ? 0 : val)
 *   this.left = (left===undefined ? null : left)
 *   this.right = (right===undefined ? null : right)
 * }
 */
/**
 * @param {TreeNode} root
 * @param {number} from
 * @param {number} to
 * @return {TreeNode}
 */
var correctBinaryTree = function (root) {
  const dfs = root => {
    if (!root || vis.has(root.right)) {
      return null;
    }
    vis.add(root);
    root.right = dfs(root.right);
    root.left = dfs(root.left);
    return root;
  };
  const vis = new Set();
  return dfs(root);
};