2023.6.12 树节点的第k个祖先

发布时间 2023-06-12 11:36:42作者: 烤肉kr

image

可以借鉴一下求LCA问题中的倍增思想。
fa[i][j]表示i号节点的第\(2^j\)个祖先。我们只需要用动态规划预处理出这个fa数组即可。
求第k个祖先,可以将k用二进制拼凑的方法划分成若干个2的整数次幂,然后利用fa数组对应地进行若干次跳跃即可,单个询问的时间复杂度\(O(logn)\)

这里由于\(n < 5 \times 10^4\),所以\(M=log(n) + 1 = 16\)

struct TreeAncestor {
    fa: Vec<Vec<i32>>
}

impl TreeAncestor {

    fn new(n: i32, parent: Vec<i32>) -> Self {
        let n = n as usize;
        let mut fa = vec![vec![-1; 16]; n];

        for i in 0..n { fa[i][0] = parent[i]; }
        for i in 0..n {
            for j in 1..16 {
                if fa[i][j - 1] != -1 { fa[i][j] = fa[fa[i][j - 1] as usize][j - 1]; }
            }
        }
        Self { fa }
    }
    
    fn get_kth_ancestor(&self, node: i32, k: i32) -> i32 {
        let mut res = node;
        for i in 0..16usize {
            if ((k >> i) & 1) == 1 { res = self.fa[res as usize][i] }
            if res == -1 { return -1; }
        }
        res
    }
}