LeetCode1483 树节点的第K个祖先——构造缓存，查询优化到O(logN)

题目

LeetCode1483 树节点的第K个祖先

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

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

树节点的第 k 个祖先节点是从该节点到根节点路径上的第 k 个节点。
提示：

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

思路

记录每个节点的第1个祖先，第10个祖先，第100个祖先，第1000个祖先和第10000个祖先，由于数据规模是10^4，因此维护到这10000就好了，用这个缓存表去查找，查找的时间复杂度是O(logN)。构造缓存的时间复杂度是O(N)

代码

class TreeAncestor {
    private int[][] anc;

    public TreeAncestor(int n, int[] parent) {
        anc = new int[n][5];
        for(int i = 0; i < n; i++) {
            Arrays.fill(anc[i], -1);
            anc[i][0] = parent[i];
        }
        for(int i = 1; i < 5; i++) {
            for(int j = 0; j < n; j++) {
                int cur = anc[j][i-1];
                for(int k = 0; k < 9; k++) if(cur != -1) {
                    cur = anc[cur][i-1];
                }
                anc[j][i] = cur;
            }
        }
    }
    
    public int getKthAncestor(int node, int k) {
        for(int i = 4; i >= 0; i--) {
            int base = 1;
            for(int j = 0; j < i; j++) {
                base *= 10;
            }
            while(k >= base) {
                node = anc[node][i];
                k -= base;
                if(node == -1) {
                    return -1;
                }
            }
        }
        return node;
    }
}