[iAlgo Insight - 树上 K 祖先] 倍增解决

Problem: 1483. 树节点的第 K 个祖先

  • Tag: Tree
  • Difficulty: Hard
  • Classic: 🌙

思路

Insight:

  • 在看到祖先问题时,都可以考虑用 二进制(倍增) 的思想解决。

考虑在面试过程中应该如何思考,我们应该如何快速找到正确的思路呢?我们来分析已有信息:

  • 已知给定的是一个 parent 数组 => 这和通常给定一个 root 节点不同,这其实告诉我们不一定需要建树(虽然我后面的代码仍然建树了)。
  • 看到 k 祖先,肯定不能暴力的去一个个往上找,在最坏的情况下是需要 \(\mathcal{O}(n)\) 的时间复杂度 => 考虑优化这个跳转过程,考虑一个例子 k = 9,我们可以转化为二进制 9 = 1001_2,如果我们能提前知道每个节点跳转 1, 2, 4, 8 ... 后的祖先是什么。那么 k = 9 可以拆分为跳转两次(分布是跳 18)。所以最后要使用倍增的思想去解决这个问题。

在第一次遇到这类问题时做不出来是正常的(但是如果做过最近公共祖先问题后仍然不会就值得反思)。但是倍增的思想我认为并不超出面试算法的查考范围,并且代码量不会很大。所以这属于一个面试好题,值得一做补充武器库。

解题方法

理解了上述思路后,问题变成了如何维护跳转数组 fa[node][binary]。也就是 node 的第 \(2^{binary}\) 个祖先。这非常简单:

  • Base case: node 的第1个祖先就是其父亲节点 father
  • Update: fa[node][binary] = fa[fa[node][binary - 1]][binary - 1]

理解了 Update 这个公式一切问题便迎刃而解。不妨假设 fa 数组变成了一个函数:

\[fa(node, binary) = node \times 2^{binary} \]

那么,Update 这个式子便可以理解为:(只需要将 \(\times\) 理解为向上跳转即可

\[\begin{aligned} fa(node, binary) &= node\times 2^{binary} \\ fa(fa(node,binary-1),binary-1) &= fa(node \times 2^{binary-1}, binary -1) \\ &= node \times 2^{binary - 1} \times 2^{binary - 1} \\ &= node \times 2^{binary} \end{aligned} \]

复杂度

时间复杂度: 二进制最多 32 位,所以是常数级别:

  • 构造 fa 数组的时间复杂度为: \(\mathcal{O}(n)\).
  • 查询的时间复杂度是 \(O(1)\) (做常数次跳转)。

空间复杂度: \(\mathcal{O}(n)\)

Code

区别于大部分题解中直接利用 parent 数组做转移,我用的是构图 + dfs 。在此题上我认为前者更好。但是在更多更复杂的情况下,构图是必要的,可以维护一些更复杂的信息(比如在最近公共祖先问题中)。

constexpr int N = 5e4 + 50;
int fa[N][31];
std::vector<int> E[N];

class TreeAncestor {
public:
    void dfs(int root, int father) {
        fa[root][0] = father;

        for (int i = 1; i < 31; ++ i) {
            // 2^{i-1} + 2^{i-1} => 2^{i}
            if (fa[root][i - 1] == -1) {
                fa[root][i] = -1;
            } else {
                fa[root][i] = fa[fa[root][i - 1]][i - 1];
            }
        }

        for (auto& go: E[root]) {
            if (go == father) continue;
            dfs(go, root);
        }
    }

    TreeAncestor(int n, vector<int>& parent) {
        memset(fa, 0, sizeof(fa));
        for (int i = 0; i < n; ++ i)
            E[i].clear();
        for (int i = 0; i < n; ++ i) {
            if (parent[i] == -1) continue;
            E[parent[i]].emplace_back(i);
        }
        
        dfs(0, -1);
    }

    int getKthAncestor(int node, int k) {
        for (int i = 0; k; ++ i, k >>= 1) {
            // printf("node: %d, curK: %d, fa[%d][%d]: %d\n", node, k, node, i, fa[node][i]);
            if (k & 1) node = fa[node][i];
            if (node == -1) return node;
        }
        return node;
    }
};

/**
 * Your TreeAncestor object will be instantiated and called as such:
 * TreeAncestor* obj = new TreeAncestor(n, parent);
 * int param_1 = obj->getKthAncestor(node,k);
 */
class TreeAncestor:

    def __init__(self, n: int, parent: List[int]):
        self.fa = [[0 for i in range(31)] for j in range(n)]
        self.E = [[] for i in range(n)]
        for i, pa in enumerate(parent):
            if pa == -1:
                continue
            self.E[pa].append(i)

        self.dfs(root=0, father=-1)

    def dfs(self, root, father):
        self.fa[root][0] = father
        for i in range(1, 31, 1):
            self.fa[root][i] = (
                -1
                if self.fa[root][i - 1] == -1
                else self.fa[self.fa[root][i - 1]][i - 1]
            )

        for go in self.E[root]:
            if go == father:
                continue
            self.dfs(go, root)

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

Reference

posted @ 2024-04-06 16:05  Last_Whisper  阅读(7)  评论(0编辑  收藏  举报