「LeetCode」1483. 树节点的第 K 个祖先——Binary Lifting

343 阅读3分钟

这是我参与11月更文挑战的第6天,活动详情查看2021最后一次更文挑战

前言

看到这个题是在知乎的一个回答里计算机基础知识对程序员来说有多重要?,这个回答提到的例子非常实际,告诉我们要重视编程中的思考、技巧和深度,重视计算机基础知识、深入了解语言特性,以此来从根本上改进我们做事的效率。

提到的这道题是LeetCode上的一道hard题,但是它最优解的代码量不多、解法的思想却非常实际,日常工作中也非常容易遇到(观察到的一个特点:比起easy、medium题来说,LeetCode上的hard题会更贴合工作实际来出)。

正文

1483. 树节点的第 K 个祖先

解析

题目描述很简单,就是给一个parent数组,从这个parent数组中我们可以构建出一整颗树,然后每次调用getKthAncestor获取某个节点的第K个祖先。

最简单的写法是不做预处理,在getKthAncestor时去递归查询具体的第K个祖先,但是这样会超时:

超时的暴力解法

class TreeAncestor:
    def __init__(self, n: int, parent: List[int]):
        self.parent = parent

    def getKthAncestor(self, node: int, k: int) -> int:
        while node != -1 and --k:
            node = self.parent[node]
            k -= 1
        return node

预处理出所有祖先节点

引入预处理的思想之后,我们可以在树的每个节点上都存储一个其所有祖先节点的数组,然后每次查询时只需要在这个数组里面查一次即可。但是这样会使我们的空间复杂度在极端情况下非常高(O(n^2)),而且在构建存储全部节点的所有祖先节点数组时需要额外的处理(因为只有parent数组即节点对其父节点的引用,构建时需要先处理出所有节点的子节点,再从根节点出发进行深度/广度遍历才能得到这个所有祖先节点的数组)。

这样的写法在查询getKthAncestor的时间开销上比起后续的Binary Lifting会更有优势,而且可能由于LeetCode数据集并不够完善,这种写法通过一定手段的处理后依然能通过数据集的测试并且时间开销非常优秀:

class TreeAncestor:
    def __init__(self, n: int, parent: List[int]):
        node2list = [None] * n
        leaves = {i for i in range(n)} - set(parent)
        for i in leaves:
            lst = []
            while i > 0:
                node2list[i] = lst, len(lst)
                i = parent[i]
                lst.append(i)
        self.__node2list = node2list

    def getKthAncestor(self, node: int, k: int) -> int:
        if node == 0:
            return -1
        lst, start = self.__node2list[node]
        index = start + k - 1
        if index < len(lst):
            return lst[index]
        return -1

可以看到,它先通过{i for i in range(n)} - set(parent)找出所有叶子节点,然后每个叶子节点创建一个路径数组保存从叶子节点往上路径上的每个节点。这样就在语言层面上节约了部分空间开销——所有parent只保存叶子节点上路径数组的引用。只是后续在index = start + k - 1需要注意计算。

Binary Lifting

这里就需要引入Binary Lifting,其本质是DP算法,即对树上每个节点都保存一个数组,其存储距某个节点距离为2^i次方的节点,只存储2^i次方节点的好处非常多:

  1. 节约空间开销;
  2. 方便parent数组的构建;
  3. 便于查询时使用位运算; 具体实现如下:
class TreeAncestor:
    def __init__(self, n: int, parent: List[int]):
        self.dp = [[-1] * 16 for _ in range(0, n)]
        for i in range(0, len(parent)):
            self.dp[i][0] = parent[i]
        for i in range(0, len(parent)):
            j = 1
            while j < 16 and self.dp[i][j-1] != -1:
                self.dp[i][j] = self.dp[self.dp[i][j-1]][j-1]
                j += 1

    def getKthAncestor(self, node: int, k: int) -> int:
        i = 15
        while i >= 0:
            if k & (1 << i):
                node = self.dp[node][i]
                if node == -1:
                    break
            i -= 1
        return node

后续会补充一点关于实现细节的解析:TODO

后记

Binary Lifting的思想在其他很多地方都有用到,比如在做236. 二叉树的最近公共祖先时可以用其把查询时间从O(n)级别优化到O(logn)级别。

在其他比较高级的算法中也用到了Binary Lifting的思想,这个后续再做介绍。