1483.Kth Ancestor of a Tree Node

Tags: Hard Dynamic Programming

Links: https://leetcode.com/problems/kth-ancestor-of-a-tree-node/

---

You are given a tree with n nodes numbered from 0 to n-1 in the form of a parent array where parent[i] is the parent of node i. The root of the tree is node 0.

Implement the function getKthAncestor (int node, int k) to return the k-th ancestor of the given node. If there is no such ancestor, return -1.

The k-th ancestor of a tree node is the k-th node in the path from that node to the root.

Example:

Input:
["TreeAncestor","getKthAncestor","getKthAncestor","getKthAncestor"]
[[7,[-1,0,0,1,1,2,2]],[3,1],[5,2],[6,3]]

Output:
[null,1,0,-1]

Explanation:
TreeAncestor treeAncestor = new TreeAncestor(7, [-1, 0, 0, 1, 1, 2, 2]);

treeAncestor.getKthAncestor(3, 1);  // returns 1 which is the parent of 3
treeAncestor.getKthAncestor(5, 2);  // returns 0 which is the grandparent of 5
treeAncestor.getKthAncestor(6, 3);  // returns -1 because there is no such ancestor

Constraints:

• 1 <= k <= n <= 5*10^4
• parent[0] == -1 indicating that 0 is the root node.
• 0 <= parent[i] < n for all 0 < i < n
• 0 <= node < n
• There will be at most 5*10^4 queries.

---

class TreeAncestor {
    vector<vector<int>> ancestor;
public:
    TreeAncestor(int n, vector<int>& parent) {
        std::ios_base::sync_with_stdio(false);
        cin.tie(NULL);
        cout.tie(NULL);

        ancestor.resize(50005, vector<int>(21));
        for (int i = 0; i < n; ++i) ancestor[i][0] = parent[i];
        for (int j = 1; j <= 20; ++j) {
            for (int i = 0; i < n; ++i) {
                ancestor[i][j] = (ancestor[i][j - 1] >= 0) ? ancestor[ancestor[i][j - 1]][j - 1] : -1;
            }
        }
    }

    int getKthAncestor(int node, int k) {
        for (int i = 0; i <= 20; ++i) {
            if (k & (1 << i)) node = ancestor[node][i];
            if (node == -1) return -1;
        }

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

倍增和二进制的思路。类似于ST表，我们用ancestor[i][j]表示节点i的第$2^j$个祖先，很显然当j = 0时，ancestor[i][0]就是parent[i]。

当j >= 1时，ancestor的状态转移方程：

ancestor[i][j] = (ancestor[i][j - 1] >= 0) ? ancestor[ancestor[i][j - 1]][j - 1] : -1;

也就是节点i的 $2^j$个祖先， 就是节点i的第$2^{j-1}$个祖先节点的第$2^{j-1}$个祖先。

因为数据范围节点数不会超过50000，所以第一个维度开50005的大小，第二个维度是21是因为数据范围到了$10^6$，一般用ST表求RMQ顶多到$10^6$的数据范围，$\log _2 10^6 \approx 19.93$，开21或25的大小都可以。

考虑getKthAncestor，比如我们想查询节点i的第13个祖先，离13最接近的是8，然后将5拆成4和1，这其实就是13的二进制表示，也就是如果想访问第13个祖先，那么先利用ancestor查找第8个祖先，将第8个祖先作为当前节点，然后查找当前节点的第4个祖先，依此类推。

最坏的情况是单链表的情形，树的深度最大是$5 \times 10^4$，查询次数是$5\times 10^4$，时间复杂度$O(n \log n)$。