题目描述

给你一棵树，树上有 n 个结点，按从 0 到 n-1 编号。树以父结点数组的形式给出，其中 parent[i] 是结点 i 的父结点。树的根结点是编号为 0 的结点。

请你设计并实现 getKthAncestor(int node, int k) 函数，函数返回结点 node 的第 k 个祖先结点。如果不存在这样的祖先结点，返回 -1。

树结点的第 k 个祖先结点是从该结点到根结点路径上的第 k 个结点。

样例

输入：
["TreeAncestor","getKthAncestor","getKthAncestor","getKthAncestor"]
[[7,[-1,0,0,1,1,2,2]],[3,1],[5,2],[6,3]]

输出：
[null,1,0,-1]

解释：
TreeAncestor treeAncestor = new TreeAncestor(7, [-1, 0, 0, 1, 1, 2, 2]);

treeAncestor.getKthAncestor(3, 1);  // 返回 1 ，它是 3 的父结点
treeAncestor.getKthAncestor(5, 2);  // 返回 0 ，它是 5 的祖父结点
treeAncestor.getKthAncestor(6, 3);  // 返回 -1 因为不存在满足要求的祖先结点

限制

• 1 <= k <= n <= 5*10^4
• parent[0] == -1 表示编号为 0 的结点是根结点。
• 对于所有的 0 < i < n，0 <= parent[i] < n 总成立。
• 0 <= node < n
• 至多有 5*10^4 次查询。

算法

(倍增) 预处理 $O(n \log n)$；查询 $O(\log n)$。

1. 预处理出每个结点 $i$ 第 2 的 $j$ 次方的祖先 $f(i, j)$。
2. 初始时，$f(i, 0) = parent(i)$。
3. 通过迭代，我们有 $f(i, j) = f(f(i, j - 1), j - 1)$。
4. 对于查询某个结点第 $k$ 个祖先，可以将 $k$ 分解为二进制的形式，根据其二进制位上 1 的位置进行迭代。

时间复杂度

• 预处理时共有 $O(\log n)$ 层，每一层需要遍历所有结点，故需要 $O(n \log n)$ 的时间。
• 每次查询只需要将 $k$ 进行二进制分解，所以查询的时间复杂度为 $O(\log n)$。

空间复杂度

• 需要额外 $O(n \log n)$ 的空间存储祖先信息。

C++ 代码

class TreeAncestor {
public:
    vector<vector<int>> f;

    TreeAncestor(int n, vector<int>& parent) {
        f.resize(n, vector<int>(16, -1));

        for (int i = 0; i < n; i++)
            f[i][0] = parent[i];

        for (int j = 1; j < 16; j++)
            for (int i = 0; i < n; i++)
                if (f[i][j - 1] != -1)
                    f[i][j] = f[f[i][j - 1]][j - 1];
    }

    int getKthAncestor(int node, int k) {
        int res = node;
        for (int i = 0; i < 16; i++)
            if ((k >> i) & 1) {
                res = f[res][i];
                if (res == -1)
                    return -1;
            }

        return res;
    }
};

/**
 * Your TreeAncestor object will be instantiated and called as such:
 * TreeAncestor* obj = new TreeAncestor(n, parent);
 * int param_1 = obj->getKthAncestor(node,k);
 */