LeetCode 1261. 在受污染的二叉树中查找元素

2022-04-03 HGNULB 约 6 分钟

文章信息 6 个分类 6 个标签

分类

题目描述

✅ 1261. 在受污染的二叉树中查找元素

题意分析

给一棵所有节点值都被改成 -1 的二叉树，但它原本的值遵循一条确定规则：根为 0，任意值为 x 的节点，其左孩子为 $2x+1$、右孩子为 $2x+2$。要求实现一个类，构造时把树还原，之后支持多次 find(target) 查询某个值是否在树中。

这是一道设计题，考察点不是单次计算而是构造与查询的代价分配。find 会被调用很多次（题目给出的调用总数上限是 $10^4$），所以应该把重活放在构造阶段，让 find 尽可能便宜。

规则本身是完全确定的：一个节点的值只由它在树中的位置决定，与树的形状无关。换句话说，只要知道从根往下的左右走法，值就唯一确定；反过来，值也唯一确定了走法。这是一个双向的一一对应。

但要注意「树中存在这个值」不等于「这个值符合编号规则」——树可能是不完整的，某个位置的编号合法但那个节点根本不存在。所以 find 必须真的检查节点是否存在，不能只验证编号的合法性。

数据规模：节点数不超过 $10^4$，值上限约 $10^6$，都不大。构造阶段做一次完整遍历完全可以接受。

边界：树只有根节点、树严重退化成一条链、查询一个不存在但编号合法的值、查询 0（根一定存在）、查询负数。

解法：DFS 复原 + 哈希集合

核心思路

构造只执行一次，而 find 最多调用上万次，因此适合在构造阶段遍历整棵树：恢复每个实际存在节点的值，并把这些值放进哈希集合，之后查询只需一次集合查找。

定义 recover(node, value)：value 是 node 按题目规则恢复后的真实值。处理当前节点后，若左孩子存在，递归传入 2*value+1；若右孩子存在，递归传入 2*value+2。

恢复不变量：每次进入 recover 时，参数 value 恰好等于当前节点由根到该位置唯一确定的编号。 根以 0 开始；若父节点编号正确，两条递推式就分别给出左右孩子的正确编号，因此可由树深归纳到全部节点。

正确性说明：每个实际节点都被访问一次并加入集合，所以集合中的值与树中恢复后的值一一对应。由此，find(target) 命中当且仅当对应位置确实存在，而不仅仅是编号在数学上合法。

编号随深度指数增长，必须核对整数范围。按根深度为 0 计算，深度 20 的最大全右链编号为 $2^{21}-2=2{,}097{,}150$，Java 的 int 与题目运行环境中的 Go int 都足够；代码也只为真实存在的孩子计算编号。若推广到更深的树，应把恢复值和集合键一起改为 64 位，而不能只扩大乘法临时变量。

解题步骤

初始化保存恢复值的哈希集合。

根非空时调用 recover(root, 0)。

在 recover 中写回当前节点值，并将它加入集合。

仅对存在的孩子计算 2*value+1 或 2*value+2 并递归。

find(target) 直接返回集合是否包含 target。

对只有右孩子的树 [-1,null,-1]，恢复后根为 0、右孩子为 2，集合为 {0,2}。find(2) 返回 true，而 find(1) 返回 false：编号 1 合法，但树中没有这个节点。

单节点树只写入 0；空孩子不会触发编号计算。稀疏树同样只记录实际存在的位置，不要求树是完全二叉树。

代码实现

import java.util.HashSet;
import java.util.Set;

class FindElements {
    private final Set<Integer> values = new HashSet<>();

    public FindElements(TreeNode root) {
        if (root != null) {
            recover(root, 0);
        }
    }

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

    private void recover(TreeNode node, int value) {
        node.val = value;
        values.add(value);

        if (node.left != null) {
            recover(node.left, value * 2 + 1);
        }
        if (node.right != null) {
            recover(node.right, value * 2 + 2);
        }
    }
}

type FindElements struct {
	values map[int]struct{}
}

func Constructor(root *TreeNode) FindElements {
	elements := FindElements{values: make(map[int]struct{})}
	if root != nil {
		elements.recover(root, 0)
	}
	return elements
}

func (elements *FindElements) Find(target int) bool {
	_, exists := elements.values[target]
	return exists
}

func (elements *FindElements) recover(node *TreeNode, value int) {
	node.Val = value
	elements.values[value] = struct{}{}

	if node.Left != nil {
		elements.recover(node.Left, value*2+1)
	}
	if node.Right != nil {
		elements.recover(node.Right, value*2+2)
	}
}

复杂度分析

设树有 n 个节点、高度为 h。

构造时间： $O(n)$，每个节点访问并插入集合一次。

查询时间： 平均 $O(1)$，来自哈希集合查找。

空间复杂度： $O(n+h)$，集合保存全部节点值，递归栈深度为 h；因 h<=n，整体为 $O(n)$。

关键点总结

构造一次、查询多次，适合用一次遍历和 $O(n)$ 存储换取平均 $O(1)$ 查询。

递归参数表示“当前节点应恢复成的值”，根成立、子节点可由父节点推出，形成完整归纳证明。

哈希集合只记录实际访问到的节点，能区分“编号合法”和“节点存在”。

恢复公式指数增长；本题树高保证 int 安全，推广时必须同步扩大节点值、集合键和计算类型。

先判断孩子存在再计算其编号，避免为不存在的位置做无意义运算。

易错点总结

把左右公式写反： 对 [-1,null,-1] 会把右孩子恢复成 1，导致 find(1) 与 find(2) 都返回错误结果。

只判断 target >= 0： 合法编号不代表节点存在，稀疏树中的空位置不能返回 true。

所有递归都继续传父节点值： 整棵树会被写成 0，集合也只剩一个值。

假设树是完全二叉树，用 target < n 判断： 两个节点的稀疏树也可能拥有编号 2，却不存在编号 1。

无视整数上界照搬到任意深度： 超过整数位宽后，孩子编号会溢出并在集合中产生碰撞；更深变体必须整体改用 64 位。

相似题目

题目	难度	考察点
919. 完全二叉树插入器	中等	同为树上设计题，但状态要支持增量插入而非一次性还原
222. 完全二叉树的节点个数	中等	直接利用编号的位结构做二分定位，不做完整遍历
297. 二叉树的序列化与反序列化	困难	结构必须显式编码，空节点也要占位，不能靠编号规则恢复
208. 实现 Trie (前缀树)	中等	同样是构造重、查询轻的设计题，但节点按字符分支而非编号