LeetCode-1483. 树节点的第 K 个祖先

原题链接：LeetCode-1483. 树节点的第 K 个祖先

题面

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

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

实现 TreeAncestor 类：

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

方法一

$pa[x][i]$代表着第$x$节点的第$2^i$个父节点的序号，即$x$的$2^{i+1}$个父节点等于$x$的第$2^i$个父节点的第$2^i$个父节点，以此类推。
$$pa[x][i+1]=pa[pa[x][i]][i]$$

class TreeAncestor {
    private int[][] pa; // pa[x][i]: x节点的2^i个父节点


    public TreeAncestor(int n, int[] parent) {
        int m = 32 - Integer.numberOfLeadingZeros(n); // 二进制长度
        pa = new int[n][m];

        // init
        for (int i = 0; i < n; i++) {
            pa[i][0] = parent[i];
        }

        for (int i = 0; i < m - 1; i++) {
            for (int x = 0; x < n; x++) {
                int p = pa[x][i];
                pa[x][i + 1] = p < 0 ? -1 : pa[p][i];
            }
        }
    }

    public int getKthAncestor(int node, int k) {
        int m = 32 - Integer.numberOfLeadingZeros(k);
        for (int i = 0; i < m; i++) {
            if (((k >> i) & 1) > 0) {
                node = pa[node][i];
                if (node < 0) break;
            }
        }
        return node;
    }

    public int getKthAncestor2(int node, int k) {
        for (; k > 0 && node != -1; k &= k - 1) {
            node = pa[node][Integer.numberOfTrailingZeros(k)];
        }
        return node;
    }
}

LCA模版

求树上任意两个节点的最近公共祖先模版方法

1. 使用DFS初始化depth数组
2. 我们可以先把更靠下的 $y$ 更新为 $y$ 的第 depth[y]-depth[x] 个祖先节点，这样 $x$ 和 $y$ 就处在同一深度了
3. 再不断从x，y往上跳跃，直到无法跳跃或x==y

循环结束，$x$和$lca$只有一步之遥，即$lca=pa[x][0]$

/**
 * LCA模版
 */
class TreeAncestor {

    private int[] depth;
    private int[][] pa;

    public TreeAncestor2(int[][] edges) {
        int n = edges.length;
        int m = 32 - Integer.numberOfLeadingZeros(n);

        List<Integer>[] g = new ArrayList[n];
        Arrays.setAll(g, e -> new ArrayList<>());

        for (var e : edges) {
            int x = e[0], y = e[1];
            g[x].add(y);
            g[y].add(x);
        }

        depth = new int[n];
        pa = new int[n][m];
        dfs(g, 0, -1);

        for (int i = 0; i < m - 1; i++) {
            for (int x = 0; x < n; x++) {
                int p = pa[x][i];
                pa[x][i + 1] = p < 0 ? -1 : pa[p][i];
            }
        }
    }

    private void dfs(List<Integer>[] g, int x, int fa) {
        pa[x][0] = fa;
        for (int y : g[x]) {
            if (y != fa) {
                depth[y] = depth[x] + 1;
                dfs(g, y, x);
            }
        }
    }


    public int getKthAncestor(int node, int k) {
        for (; k > 0; k &= k - 1) {
            node = pa[node][Integer.numberOfTrailingZeros(k)];
        }
        return node;
    }

    public int getLCA(int x, int y) {
        if (depth[x] > depth[y]) {
            int t = y;
            y = x;
            x = t;
        }

        //使得x,y在同一个深度
        y = getKthAncestor(y, depth[y] - depth[x]);

        if (y == x) {
            return x;
        }

        for (int i = pa[x].length - 1; i >= 0; i--) {
            int px = pa[x][i], py = pa[y][i];
            if (px != py) {
                x = px;
                y = py;
            }
        }

        return pa[x][0];
    }
}

相关题目

2846. 边权重均等查询

参考资料

https://leetcode.cn/problems/kth-ancestor-of-a-tree-node/solutions/2305895/mo-ban-jiang-jie-shu-shang-bei-zeng-suan-v3rw