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 。

说明:

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

思路

这个题让我们维护一个数据结构,来查找树中任意节点的第k个祖先节点。直接的想法是保存每一个节点的父节点,需要的时候直接根据下标获取。刚开始用的 int[][] 超出了空间限制,后来改成 List<Integer>[] 虽然多通过了几个测试用例,但是后面会超时。仔细分析最坏的情况下(所有节点仅有一个子树的情况),需要添加 n(n+1)/2 个父节点(首项为1,公差为1的等差数列求和),时间复杂度是O(n^2)。

一个解决办法是不要保存重复的父节点,以只有一个子树的情况举例,最后一个节点第k个祖先,就是其父节点的第k-1个祖先。如果这个节点已经保存有祖先节点的信息,就无需重复计算了。

所以我的解决方案就是使用缓存,如果父节点的祖先信息没有保存,就将当前节点的祖先信息写入缓存,直到遇到存在缓存的祖先节点,如果它记录的祖先节点个数大于k - cnt就直接返回,否则继续向该缓存的祖先节点集合添加,直到遇到下一个有缓存的节点或者cnt == k

这种方法虽然能够通过,但是与测试用例的的顺序是有关的,如果是从子节点逐步向前测试的话,缓存一直不命中,时间复杂度还是O(n^2)。

官方的解法使用的是倍增的思想,好像还挺常用的,算是个模板算法。核心思想是保存当前节点的父节点,爷爷节点,爷爷的爷爷节点......,即每个节点 x 的第 2^i 个祖先节点。这样不论k取什么值,都可以分解为不同的2的幂之和,然后向前查找即可。预处理的时间复杂度是O(nlogn),查询的时间复杂度是O(logk)。

代码

/**
 * @date 2024-04-06 9:45
 */
public class TreeAncestor1483 {

    /**倍增的写法 */
    public static class TreeAncestor_v4 {

        int[][] dp;

        public TreeAncestor_v4(int n, int[] parent) {
            dp = new int[16][];
            dp[0] = parent;
            for (int i = 1; i < 16; i++) {
                dp[i] = new int[n];
                Arrays.fill(dp[i], -1);
            }

            for (int i = 1; i < 16; i++) {
                for (int j = 0; j < n; j++) {
                    if (dp[i - 1][j] != -1) {
                        dp[i][j] = dp[i - 1][dp[i - 1][j]];
                    }
                }
            }
        }

        public int getKthAncestor(int node, int k) {
            int p = node;
            int b = 0;
            int mod;
            while (k != 0) {
                mod = k & 1;
                if (mod == 1) {
                    p = dp[b][p];
                    if (p == -1) {
                        return -1;
                    }
                }
                k = k >> 1;
                b++;
            }
            return p;
        }
    }

    int[] parent;
    List<Integer>[] cache;

    public TreeAncestor1483(int n, int[] parent) {
        this.parent = parent;
        cache = new ArrayList[n];
        for (int i = 0; i < cache.length; i++) {
            cache[i] = new ArrayList<>();
        }
    }

    public int getKthAncestor(int node, int k) {
        if (node == -1) {
            return -1;
        }
        int cnt = 0;
        int p = node;
        while (cnt != k && p != -1) {
            if (cache[p].size() == 0) {
                cache[node].add(parent[p]);
                p = parent[p];
                cnt++;
            } else {
                if (cache[p].size() >= k - cnt) {
                    return cache[p].get(k - cnt - 1);
                } else {
                    cnt += cache[p].size();
                    node = p;
                    p = cache[p].get(cache[p].size() - 1);
                }
            }

        }
        return p;
    }
}

性能

这里是使用缓存写法的耗时,官方题解的耗时差不多也是这个样。

使用倍增的写法

发表回复

您的邮箱地址不会被公开。 必填项已用 * 标注