[iAlgo Insight - 树上 K 祖先] 倍增解决
Problem: 1483. 树节点的第 K 个祖先
- Tag:
Tree - Difficulty:
Hard - Classic: 🌙
思路
Insight:
- 在看到祖先问题时，都可以考虑用 二进制（倍增) 的思想解决。
考虑在面试过程中应该如何思考，我们应该如何快速找到正确的思路呢？我们来分析已有信息：
- 已知给定的是一个
parent数组=>这和通常给定一个root节点不同，这其实告诉我们不一定需要建树（虽然我后面的代码仍然建树了）。 - 看到
k祖先，肯定不能暴力的去一个个往上找，在最坏的情况下是需要 \(\mathcal{O}(n)\) 的时间复杂度=>考虑优化这个跳转过程，考虑一个例子k = 9，我们可以转化为二进制9 = 1001_2，如果我们能提前知道每个节点跳转1,2,4,8... 后的祖先是什么。那么k = 9可以拆分为跳转两次（分布是跳1和8）。所以最后要使用倍增的思想去解决这个问题。
在第一次遇到这类问题时做不出来是正常的（但是如果做过最近公共祖先问题后仍然不会就值得反思）。但是倍增的思想我认为并不超出面试算法的查考范围，并且代码量不会很大。所以这属于一个面试好题，值得一做补充武器库。
解题方法
理解了上述思路后，问题变成了如何维护跳转数组 fa[node][binary]。也就是 node 的第 \(2^{binary}\) 个祖先。这非常简单：
Base case:node的第1个祖先就是其父亲节点father。Update:fa[node][binary] = fa[fa[node][binary - 1]][binary - 1]。
理解了 Update 这个公式一切问题便迎刃而解。不妨假设 fa 数组变成了一个函数:
\[fa(node, binary) = node \times 2^{binary}
\]
那么，Update 这个式子便可以理解为：（只需要将 \(\times\) 理解为向上跳转即可）
\[\begin{aligned}
fa(node, binary) &= node\times 2^{binary} \\
fa(fa(node,binary-1),binary-1) &= fa(node \times 2^{binary-1}, binary -1) \\
&= node \times 2^{binary - 1} \times 2^{binary - 1} \\
&= node \times 2^{binary}
\end{aligned}
\]
复杂度
时间复杂度: 二进制最多 32 位，所以是常数级别：
- 构造
fa数组的时间复杂度为： \(\mathcal{O}(n)\). - 查询的时间复杂度是 \(O(1)\) （做常数次跳转）。
空间复杂度: \(\mathcal{O}(n)\)
Code
区别于大部分题解中直接利用 parent 数组做转移，我用的是构图 + dfs 。在此题上我认为前者更好。但是在更多更复杂的情况下，构图是必要的，可以维护一些更复杂的信息（比如在最近公共祖先问题中）。
constexpr int N = 5e4 + 50;
int fa[N][31];
std::vector<int> E[N];
class TreeAncestor {
public:
void dfs(int root, int father) {
fa[root][0] = father;
for (int i = 1; i < 31; ++ i) {
// 2^{i-1} + 2^{i-1} => 2^{i}
if (fa[root][i - 1] == -1) {
fa[root][i] = -1;
} else {
fa[root][i] = fa[fa[root][i - 1]][i - 1];
}
}
for (auto& go: E[root]) {
if (go == father) continue;
dfs(go, root);
}
}
TreeAncestor(int n, vector<int>& parent) {
memset(fa, 0, sizeof(fa));
for (int i = 0; i < n; ++ i)
E[i].clear();
for (int i = 0; i < n; ++ i) {
if (parent[i] == -1) continue;
E[parent[i]].emplace_back(i);
}
dfs(0, -1);
}
int getKthAncestor(int node, int k) {
for (int i = 0; k; ++ i, k >>= 1) {
// printf("node: %d, curK: %d, fa[%d][%d]: %d\n", node, k, node, i, fa[node][i]);
if (k & 1) node = fa[node][i];
if (node == -1) return node;
}
return node;
}
};
/**
* Your TreeAncestor object will be instantiated and called as such:
* TreeAncestor* obj = new TreeAncestor(n, parent);
* int param_1 = obj->getKthAncestor(node,k);
*/
class TreeAncestor:
def __init__(self, n: int, parent: List[int]):
self.fa = [[0 for i in range(31)] for j in range(n)]
self.E = [[] for i in range(n)]
for i, pa in enumerate(parent):
if pa == -1:
continue
self.E[pa].append(i)
self.dfs(root=0, father=-1)
def dfs(self, root, father):
self.fa[root][0] = father
for i in range(1, 31, 1):
self.fa[root][i] = (
-1
if self.fa[root][i - 1] == -1
else self.fa[self.fa[root][i - 1]][i - 1]
)
for go in self.E[root]:
if go == father:
continue
self.dfs(go, root)
def getKthAncestor(self, node: int, k: int) -> int:
for i in range(31):
if k & 1:
node = self.fa[node][i]
if node == -1:
return node
k >>= 1
return node
