https://leetcode.cn/problems/kth-ancestor-of-a-tree-node/description/?envType=daily-question&envId=2024-04-11
非常好二分,使我思路旋转
题目给了每个节点的父节点,最容易想到的就是直接暴力搜。简单算一下最坏的情况下时间复杂度为 $O(n^2)$ ,显然会超时。
考虑使用一个二维数组记忆化搜索,空间复杂度 $O(n^2)$ ,显然会爆内存。
题目要求节点的第 k 个祖先,与中间一次走了几步无关,所以可以使用倍增的思想,记录下每个节点走 $2^0, 2^1, ... , 2^j$ 步后的节点,如果走不到记为 -1 即可。题目中 n 最大为 $5 \times 10^4$ ,因此 $j$ 取 16 即可覆盖所有的情况。此时这个稀疏表空间复杂度为 $O(nlogn)$ ,处于可接受的范围内了。
有了稀疏表之后,只需要对 $k$ 的二进制从低到高遍历一遍,第 $i$ 位为 1 就代表节点应该向前走 $2^i$ 步,即取表中对应节点的第 $i$ 个元素替换当前节点。比如 $k$ 为 26 时,二进制表示为 11010,意味着需要 node 节点向前走 $2^1 + 2^3 + 2^4$ 步,所以只需要执行 node = table[node][i] (i = 1, 3, 4)
即可得到答案。
因此最终解答如下:
class TreeAncestor {
private:
const static int MAX_LOGN = 16;
int** ancestors;
public:
TreeAncestor(int n, std::vector<int>& parent) {
ancestors = new int* [n];
for (int i = 0; i < n; ++i) {
ancestors[i] = new int [MAX_LOGN];
std::fill(ancestors[i], ancestors[i] + MAX_LOGN, -1);
ancestors[i][0] = parent[i];
}
for (int i = 1; i < MAX_LOGN; ++i) {
for (int j = 0; j < n; ++j) {
if (ancestors[j][i - 1] != -1)
ancestors[j][i] = ancestors[ancestors[j][i - 1]][i - 1];
}
}
}
int getKthAncestor(int node, int k) {
for (int i = 0; i < MAX_LOGN; ++i) {
if ((k >> i) & 1) {
node = ancestors[node][i];
if (node == -1)
return -1;
}
}
return node;
}
};