树中最接近路径的节点

Submission

运行时间: 52 ms

内存: 17.8 MB

class HLD:
    def __init__(self, g, root): 
        #无论是点还是dfn还是dep,都从1开始,默认0是无
        n=len(g)
        self.g=g
        self.fa=[0]*(n+5)   #父节点,0表示无父节点
        self.size=[1]*(n+5) #子树大小
        self.dep=[0]*(n+5)  #深度,根深度为1
        self.son=[0]*(n+5)  #重儿子,0表示无儿子
        self.dfn=[0]*(n+5)  #dfs序,子树终点的dfs序是dfn[i]+size[i]-1
        self.top=[0]*(n+5)  #所在重链起点,起点就是自己
        self.rank=[0]*(n+5) #dfs序为i的节点编号

        fa=self.fa;size=self.size;dep=self.dep;son=self.son
        dfn=self.dfn;top=self.top;rank=self.rank
        stk=[[root,0,0]] #node,flag,fa
        dep[root]=1
        while stk:
            u,flag,father=stk.pop()
            if flag:
                for v in g[u]:
                    if v!=father:
                        size[u]+=size[v]
                        if son[u]==0 or size[v]>size[son[u]]:
                            son[u]=v
            else:
                stk.append([u,1,father])
                for v in g[u]:
                    if v!=father:
                        stk.append([v,0,u])
                        fa[v]=u
                        dep[v]=dep[u]+1
        stk=[[root,root]]
        tot=1
        while stk:
            u,tops=stk.pop()
            dfn[u]=tot
            rank[tot]=u
            tot+=1
            top[u]=tops
            if son[u]==0:
                continue
            for v in g[u]:
                if v!=fa[u] and v!=son[u]:
                    stk.append([v,v])
            stk.append([son[u],tops])
    def lca(self,u,v):  #求u和v的最近公共祖先节点
        fa=self.fa;size=self.size;dep=self.dep;son=self.son
        dfn=self.dfn;top=self.top;rank=self.rank
        while top[u]!=top[v]:
            if dep[top[u]]>dep[top[v]]:
                u=fa[top[u]]
            else:
                v=fa[top[v]]
        return v if dep[u]>dep[v] else u
    
    def dis(self,u,v):
        dep=self.dep
        return dep[u]+dep[v]-2*dep[self.lca(u,v)]
class Solution:
    def closestNode(self, n: int, edges: List[List[int]], query: List[List[int]]) -> List[int]:
        g=[[] for _ in range(n+1)]
        for u,v in edges:
            u+=1;v+=1
            g[u]+=[v]
            g[v]+=[u]
        hld=HLD(g,1)
        dfn=hld.dfn
        size=hld.size
        res=[]
        for u,v,x in query:
            u+=1;v+=1;x+=1
            lca=hld.lca(u,v)
            if dfn[lca]<=dfn[x]<dfn[lca]+size[lca]:
                l1=hld.lca(u,x)
                l2=hld.lca(v,x)
                if hld.dis(l1,x)<hld.dis(l2,x):
                    res.append(l1-1)
                else:
                    res.append(l2-1)
            else:
                res.append(lca-1)
        return res

Explain

该题解使用了重链剖分(Heavy-Light Decomposition, HLD)的思路来解决树中查询最接近路径的节点的问题。具体思路如下: 1. 通过DFS遍历树,计算每个节点的父节点、子树大小、深度、重儿子等信息。 2. 将树划分为若干重链,每个重链由重儿子连接而成。对于每个节点,记录其所在重链的起点、DFS序等信息。 3. 通过LCA (最近公共祖先)和重链剖分,可以快速计算树中任意两点之间的距离。 4. 对于每个查询(u,v,x),首先判断x是否在u和v的最近公共祖先lca的子树中: - 如果在,则分别计算x到u和v的LCA的距离,取较近的作为答案。 - 如果不在,则直接返回lca作为最接近路径的节点。

时间复杂度: O(n+qlogn)

空间复杂度: O(n)

class HLD:
    def __init__(self, g, root): 
        #无论是点还是dfn还是dep,都从1开始,默认0是无
        n=len(g)
        self.g=g
        self.fa=[0]*(n+5)   #父节点,0表示无父节点
        self.size=[1]*(n+5) #子树大小
        self.dep=[0]*(n+5)  #深度,根深度为1
        self.son=[0]*(n+5)  #重儿子,0表示无儿子
        self.dfn=[0]*(n+5)  #dfs序,子树终点的dfs序是dfn[i]+size[i]-1
        self.top=[0]*(n+5)  #所在重链起点,起点就是自己
        self.rank=[0]*(n+5) #dfs序为i的节点编号

        fa=self.fa;size=self.size;dep=self.dep;son=self.son
        dfn=self.dfn;top=self.top;rank=self.rank
        stk=[[root,0,0]] #node,flag,fa
        dep[root]=1
        while stk:
            u,flag,father=stk.pop()
            if flag:
                for v in g[u]:
                    if v!=father:
                        size[u]+=size[v]
                        if son[u]==0 or size[v]>size[son[u]]:
                            son[u]=v
            else:
                stk.append([u,1,father])
                for v in g[u]:
                    if v!=father:
                        stk.append([v,0,u])
                        fa[v]=u
                        dep[v]=dep[u]+1
        stk=[[root,root]]
        tot=1
        while stk:
            u,tops=stk.pop()
            dfn[u]=tot
            rank[tot]=u
            tot+=1
            top[u]=tops
            if son[u]==0:
                continue
            for v in g[u]:
                if v!=fa[u] and v!=son[u]:
                    stk.append([v,v])
            stk.append([son[u],tops])
    def lca(self,u,v):  #求u和v的最近公共祖先节点
        fa=self.fa;size=self.size;dep=self.dep;son=self.son
        dfn=self.dfn;top=self.top;rank=self.rank
        while top[u]!=top[v]:
            if dep[top[u]]>dep[top[v]]:
                u=fa[top[u]]
            else:
                v=fa[top[v]]
        return v if dep[u]>dep[v] else u
    
    def dis(self,u,v):
        dep=self.dep
        return dep[u]+dep[v]-2*dep[self.lca(u,v)]
class Solution:
    def closestNode(self, n: int, edges: List[List[int]], query: List[List[int]]) -> List[int]:
        g=[[] for _ in range(n+1)]
        for u,v in edges:
            u+=1;v+=1
            g[u]+=[v]
            g[v]+=[u]
        hld=HLD(g,1)
        dfn=hld.dfn
        size=hld.size
        res=[]
        for u,v,x in query:
            u+=1;v+=1;x+=1
            lca=hld.lca(u,v)
            if dfn[lca]<=dfn[x]<dfn[lca]+size[lca]:
                l1=hld.lca(u,x)
                l2=hld.lca(v,x)
                if hld.dis(l1,x)<hld.dis(l2,x):
                    res.append(l1-1)
                else:
                    res.append(l2-1)
            else:
                res.append(lca-1)
        return res

Explore

在构建重链剖分时,选择每个节点的重儿子基于子树大小的标准。具体地,对于每个节点,遍历其所有子节点,并根据每个子节点的子树大小选择最大的一个作为重儿子。选择最大子树大小的重儿子是为了确保剖分后的重链尽可能长,从而减少从任意节点到根节点路径上所经过的重链数量,优化查询效率。

确保DFS遍历和重链剖分过程中计算的DFS序和子树大小准确性,依赖于正确的递归遍历和数据结构的完整性。通过使用栈(或递归)来遍历整个图,并在第一次访问节点时标记该节点,并在遍历返回时更新子树大小,可以确保每个节点和连接只被访问一次。如果图的表示(例如邻接表)是完整的,并且遍历过程中没有遗漏任何节点或边,那么计算是准确的。确保没有遗漏的关键是图的构建过程中必须包含所有节点和边的正确添加。

在计算LCA时,比较`top[u]`和`top[v]`的深度而不是直接比较`u`和`v`的深度,是因为重链剖分的目的是将树分解成多个重链,每个重链可以视为一个单独的路径。通过比较重链顶点的深度,可以快速确定`u`和`v`是否在同一重链上,或者哪一个节点更接近根节点。这样做可以迅速将问题规模缩小,通过跳过整个重链来减少递归深度或迭代次数,从而快速找到`u`和`v`的最近公共祖先。

在查询时需要将节点的索引加1,确实是因为原始的节点编号是从0开始的,而在HLD类中的处理是从1开始。这种处理方式常见于算法实现中,因为在某些编程语言或习惯中,从1开始的索引可以简化一些数学运算,尤其是在树形结构和图论中。因此,在将输入数据传递给HLD类之前,需要将节点索引统一加1以适应HLD类内部从1开始的处理方式。