公共祖先问题

前言：逃不开的「公共祖先」问题

0X00 一次查询

用 set 写比较简单，注意其中一个是另外一个祖先的情况。

class Solution:
    def lowestCommonAncestor(self, root: 'TreeNode', a: 'TreeNode', b: 'TreeNode') -> 'TreeNode':
        parent = {}
        def dfs(r, p):
            if r is None: return
            parent[r] = p
            dfs(r.left, r)
            dfs(r.right, r)

        dfs(root, None)

        p, path = parent[a], set([a])
        while p:
            path.add(p)
            p = parent[p]
        
        if b in path: return b
        p = parent[b]
        while p not in path:
            p = parent[p]
        
        return p

0X01 多次查询

使用「倍增」实现 $O(logn)$ 的查询

首先通过这道题理解「倍增」的思路：

1483. 树节点的第 K 个祖先

class TreeAncestor:

    def __init__(self, n: int, parent: List[int]):
        n = len(parent)
        dp = [[-1] * 20 for _ in range(n)]
        for i in range(n): dp[i][0] = parent[i]
        # 每个点的 i-1 都先做了
        for i in range(1, 20):
            for j in range(n):
                if dp[j][i-1] != -1: dp[j][i] = dp[dp[j][i-1]][i-1]
        self.dp = dp

    def getKthAncestor(self, node: int, k: int) -> int:
        nums = []
        while k:
            nums.append(k % 2)
            k //= 2
        p, dp = node, self.dp
        for i, v in enumerate(nums):
            if v == 0: continue
            p = dp[p][i]
            if p == -1: return -1
        return p

dp[i][j] 表示节点 i 的上面 $2^j$ 个祖先。转移方程为：

dp[i][j] = dp[dp[i][j-1]][j-1] 意思是：节点 i 的上面 $2^j$ 个祖先等于节点 i 上面 $2^{j-1}$ 个祖先的 $2^{j-1}$ 个祖先。

最后只需要将 k 以 2 为底分解就能求出答案。

接着通过倍增我们来求 lca

首先我们预处理出深度，将二者深度跳到统一起跑线上
同时向上跳，直到相同

代码如下：

from math import log2
def lca(x, y):
    if dep[x] < dep[y]: return lca(x, y) # 确保 dep[x] > dep[y]
    while dep[x] > dep[y]: x = fa[x][int(log2(dep[x] - dep[y]))] # 保证高度一致
    if x == y: return x # 特判
    for i in range(int(log2(dep[x])), -1, -1): # 一起大步向上跳
        if fa[x][i] != fa[y][i]: x, y = fa[x][i], fa[y][i]
    return fa[x][0]

其中的 fa 就是上面代码的 dp

最后编辑于：2020.07.28 17:22:14