【Leetcode】1483. Kth Ancestor of a Tree Node

题目地址：

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

给定一个$n$个节点的树，每个节点编号$0\sim n-1$，另外给出每个节点的父亲节点是谁，以数组$p$给出。要求设计一个数据结构可以应答这样的询问，每次询问给出一个节点编号$u$和一个正整数$k$，问$u$的第$k$个祖先是谁（即$u$向上走$k$步是谁）。如果该祖先不存在则返回$-1$。

可以用倍增法。设$f[u][k]$是节点$u$向上走$2^k$步能走到的节点，那么有$$f[u][k]=f[f[u][k-1]][k-1]$$初始条件$f[u][0]=p[u]$，从而可以递推出$f$数组。因为一共有$n$个节点，树最高就是$n$，也就是说向上最多能走$n-1$步，我们找到最小的$2^k\ge n-1$，把$f$的第二维开$k=\lfloor {\log }_{2}n-1\rfloor +1$这么长，以保证询问都能被正确应答。接下来考虑如何应答询问，可以用类似快速幂的思想，比方说$k$的二进制表示是$1101_2$，那么$1101_2=2^0+2^2+2^3$，那么可以先求$u_1=f[u][0]$，这样就跳了$2^0$步，再接着求$u_2=f[u_1][2]$，这样又跳了$2^2$步，再接者求$f[u_2][3]$，这样就将$k$步全跳完了。中途如果发现要跳的幂次大于了$f$的第二维能取的最大值，或者跳到了$-1$，则直接返回$-1$。代码如下：

public class TreeAncestor {
    
    private int[][] p;
    private int log;
    
    public TreeAncestor(int n, int[] parent) {
        log = (int) (Math.log(n - 1) / Math.log(2)) + 1;
        p = new int[n][log];
        // 初始化p数组
        for (int i = 0; i < parent.length; i++) {
            p[i][0] = parent[i];
        }
        // 按公式递推p数组
        for (int i = 1; i < log; i++) {
            for (int j = 0; j < n; j++) {
                if (p[j][i - 1] != -1) {
                    p[j][i] = p[p[j][i - 1]][i - 1];
                } else {
                    p[j][i] = -1;
                }
            }
        }
    }
    
    public int getKthAncestor(int node, int k) {
        int pow = 0;
        while (k > 0) {
            if (pow >= log || node == -1) {
                return -1;
            }
            
            if ((k & 1) == 1) {
                node = p[node][pow];
            }
            
            k >>= 1;
            pow++;
        }
        
        return node;
    }
}

预处理时间复杂度$O(n\log n)$，每次询问时间$O(\log n)$，空间$O(n\log n)$。

C++：

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

  TreeAncestor(int n, vector<int> &p) {
    int m = (int)log2(n) + 1;
    f = vector<vector<int>>(n, vector<int>(m, -1));
    for (int i = 0; i < p.size(); i++) f[i][0] = p[i];

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

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

时空复杂度一样。