LeetCode第1261题 - 在受污染的二叉树中查找元素

题目

解答

class FindElements {
  Set<Integer> set = new HashSet<>();

  public FindElements(TreeNode root) {
    if (root == null) {
      return;
    }

    LinkedList<TreeNode> queue = new LinkedList<>();
    queue.add(root);
    root.val = 0;
    while (!queue.isEmpty()) {
      TreeNode node = queue.removeFirst();
      set.add(node.val);
      int value = node.val * 2;

      if (node.left != null) {
        node.left.val = value + 1;
        queue.add(node.left);
      }

      if (node.right != null) {
        node.right.val = value + 2;
        queue.add(node.right);
      }
    }

  }

  public boolean find(int target) {
    return set.contains(target);
  }
}

总结

使用广度优先算法遍历树，同时将节点的值保存起来，以待后续的查询操作。