LeetCode 1660.纠正二叉树(Medium)

TreeNode root
int fromNode （在 correctBinaryTree 中不可见）
int toNode （在 correctBinaryTree 中不可见）
当以 root 为根的二叉树被解析后，值为 fromNode 的节点 TreeNode 将其右子节点指向值为 toNode 的节点 TreeNode 。
然后， root 传入 correctBinaryTree 的参数中。

示例 1

输入: root = [1,2,3], fromNode = 2, toNode = 3
输出: [1,null,3]
解释: 值为 2 的节点是无效的，所以移除之。

示例 2:

输入: root = [8,3,1,7,null,9,4,2,null,null,null,5,6], fromNode = 7, toNode = 4
输出: [8,3,1,null,null,9,4,null,null,5,6]
解释: 值为 7 的节点是无效的，所以移除这个节点及其子节点 2。

/**
 * 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 {
    public TreeNode correctBinaryTree(TreeNode root) {
        Queue<TreeNode> queue = new LinkedList();
        HashSet<TreeNode> set = new HashSet<>();
        queue.offer(root);
        int size = 1;
        while (size > 0) {
            TreeNode cur = queue.poll();
            size--;
            if (cur.right != null) {
                set.add(cur.right);
                queue.offer(cur.right);
                // cur.right.right本应该再下一层，如果出现在当前层则说明这个节点的right指针连错了
                if (set.contains(cur.right.right)) {
                    cur.right = null;
                    break;
                }
            }
            if (cur.left != null) {
                set.add(cur.left);
                queue.offer(cur.left);
                // cur.left.right本应该再下一层，如果出现在当前层则说明这个节点的right指针连错了
                if (set.contains(cur.left.right)) {
                    cur.left = null;
                    break;
                }
            }
            if (size == 0) {
                size = queue.size();
                set = new HashSet<>();
            }
        }
        return root;
    }
}