https://leetcode.cn/problems/kth-ancestor-of-a-tree-node/description/?envType=daily-question&envId=2024-04-11

非常好二分,使我思路旋转

题目给了每个节点的父节点,最容易想到的就是直接暴力搜。简单算一下最坏的情况下时间复杂度为 $O(n^2)$ ,显然会超时。

考虑使用一个二维数组记忆化搜索,空间复杂度 $O(n^2)$ ,显然会爆内存。

题目要求节点的第 k 个祖先,与中间一次走了几步无关,所以可以使用倍增的思想,记录下每个节点走 $2^0, 2^1, ... , 2^j$ 步后的节点,如果走不到记为 -1 即可。题目中 n 最大为 $5 \times 10^4$ ,因此 $j$ 取 16 即可覆盖所有的情况。此时这个稀疏表空间复杂度为 $O(nlogn)$ ,处于可接受的范围内了。

有了稀疏表之后,只需要对 $k$ 的二进制从低到高遍历一遍,第 $i$ 位为 1 就代表节点应该向前走 $2^i$ 步,即取表中对应节点的第 $i$ 个元素替换当前节点。比如 $k$ 为 26 时,二进制表示为 11010,意味着需要 node 节点向前走 $2^1 + 2^3 + 2^4$ 步,所以只需要执行 node = table[node][i] (i = 1, 3, 4) 即可得到答案。

因此最终解答如下:

class TreeAncestor {
private:
    const static int MAX_LOGN = 16;
    int** ancestors;
public:
    TreeAncestor(int n, std::vector<int>& parent) {
        ancestors = new int* [n];
        for (int i = 0; i < n; ++i) {
            ancestors[i] = new int [MAX_LOGN];
            std::fill(ancestors[i], ancestors[i] + MAX_LOGN, -1);
            ancestors[i][0] = parent[i];
        }

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

    int getKthAncestor(int node, int k) {
        for (int i = 0; i < MAX_LOGN; ++i) {
            if ((k >> i) & 1) {
                node = ancestors[node][i];
                if (node == -1)
                    return -1;
            }
        }
        return node;
    }
};
最后修改:2024 年 04 月 21 日
如果觉得我的文章对你有用,请随意赞赏