[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 。在此题上我认为前者更好。但是在更多更复杂的情况下，构图是必要的，可以维护一些更复杂的信息（比如在最近公共祖先问题中）。

Copy
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);
 */

Copy
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

Reference#

• OI WIKI: 最近公共祖先
• LeetCode 236. Lowest Common Ancestor of a Binary Tree
• LeetCode 235. Lowest Common Ancestor of a Binary Search Tree