Hard

题目描述

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

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

实现 TreeAncestor 类：

  • TreeAncestor(int n, int[] parent) 使用树中节点的数量和父节点数组初始化对象。
  • int getKthAncestor(int node, int k) 返回节点 node 的第 k 个祖先节点。如果不存在这样的祖先节点，返回 -1 。

示例 1：

输入：
["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.length == n
  • parent[0] == -1
  • 对于所有的 0 < i < n ，0 <= parent[i] < n 成立
  • 0 <= node < n
  • 至多查询 5 * 10^4 次

解题思路

这道题需要高效地查询树中节点的第k个祖先。如果每次查询都从当前节点开始向上遍历k步，时间复杂度会达到O(k)，在最坏情况下可能导致超时。

核心思路是使用倍增法（Binary Lifting）：

倍增法是一种预处理技巧，我们预先计算每个节点的第2^0、2^1、2^2、…个祖先。具体来说：

  • up[i][j] 表示节点i的第2^j个祖先
  • up[i][0] = parent[i]（第1个祖先就是父节点）
  • up[i][j] = up[up[i][j-1]][j-1]（第2^j个祖先 = 第2^(j-1)个祖先的第2^(j-1)个祖先）

查询时，我们将k表示为二进制形式，然后根据k的二进制位来决定向上跳跃的步数。例如，如果k=5（二进制101），我们先跳2^2=4步，再跳2^0=1步。

时间复杂度分析：

  • 预处理：O(n log n)
  • 每次查询：O(log k)

这种方法能够在对数时间内完成查询，非常适合处理大量查询的场景。

代码实现

class TreeAncestor {
private:
    vector<vector<int>> up;
    int LOG;
    
public:
    TreeAncestor(int n, vector<int>& parent) {
        LOG = ceil(log2(n)) + 1;
        up.assign(n, vector<int>(LOG, -1));
        
        // 初始化第1个祖先（2^0 = 1）
        for (int i = 0; i < n; i++) {
            up[i][0] = parent[i];
        }
        
        // 倍增预处理
        for (int j = 1; j < LOG; j++) {
            for (int i = 0; i < n; i++) {
                if (up[i][j-1] != -1) {
                    up[i][j] = up[up[i][j-1]][j-1];
                }
            }
        }
    }
    
    int getKthAncestor(int node, int k) {
        for (int i = 0; i < LOG; i++) {
            if ((k >> i) & 1) {
                node = up[node][i];
                if (node == -1) break;
            }
        }
        return node;
    }
};

class TreeAncestor:

    def __init__(self, n: int, parent: List[int]):
        import math
        self.LOG = int(math.log2(n)) + 1
        self.up = [[-1] * self.LOG for _ in range(n)]
        
        # 初始化第1个祖先（2^0 = 1）
        for i in range(n):
            self.up[i][0] = parent[i]
        
        # 倍增预处理
        for j in range(1, self.LOG):
            for i in range(n):
                if self.up[i][j-1] != -1:
                    self.up[i][j] = self.up[self.up[i][j-1]][j-1]

    def getKthAncestor(self, node: int, k: int) -> int:
        for i in range(self.LOG):
            if (k >> i) & 1:
                node = self.up[node][i]
                if node == -1:
                    break
        return node

public class TreeAncestor {
    private int[,] up;
    private int LOG;
    
    public TreeAncestor(int n, int[] parent) {
        LOG = (int)Math.Ceiling(Math.Log2(n)) + 1;
        up = new int[n, LOG];
        
        // 初始化为-1
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < LOG; j++) {
                up[i, j] = -1;
            }
        }
        
        // 初始化第1个祖先（2^0 = 1）
        for (int i = 0; i < n; i++) {
            up[i, 0] = parent[i];
        }
        
        // 倍增预处理
        for (int j = 1; j < LOG; j++) {
            for (int i = 0; i < n; i++) {
                if (up[i, j-1] != -1) {
                    up[i, j] = up[up[i, j-1], j-1];
                }
            }
        }
    }
    
    public int GetKthAncestor(int node, int k) {
        for (int i = 0; i < LOG; i++) {
            if (((k >> i) & 1) == 1) {
                node = up[node, i];
                if (node == -1) break;
            }
        }
        return node;
    }
}

var TreeAncestor = function(n, parent) {
    this.LOG = 16;
    this.up = Array(n).fill().map(() => Array(this.LOG).fill(-1));
    
    for (let i = 0; i < n; i++) {
        this.up[i][0] = parent[i];
    }
    
    for (let j = 1; j < this.LOG; j++) {
        for (let i = 0; i < n; i++) {
            if (this.up[i][j-1] !== -1) {
                this.up[i][j] = this.up[this.up[i][j-1]][j-1];
            }
        }
    }
};

TreeAncestor.prototype.getKthAncestor = function(node, k) {
    for (let i = 0; i < this.LOG; i++) {
        if ((k >> i) & 1) {
            node = this.up[node][i];
            if (node === -1) return -1;
        }
    }
    return node;
};
在 LeetCode 上运行： 力扣中文版 LeetCode 国际版

复杂度分析

操作时间复杂度空间复杂度
初始化O(n log n)O(n log n)
查询O(log k)O(1)

相关题目