Submission #7488237

---

Source Code Expand

def resolve():
    class LCA(object):
        def __init__(self,edges,root):
            self.edges=edges
            self.root=root
            self.n=len(edges) # n=|V|
            self.logn=(self.n-1).bit_length() # logn = ceil(log2(n))
            # initialization
            self.depth=[float("inf") if i!=root else 0 for i in range(self.n)]
            self.parent=[[-1]*self.logn for _ in range(self.n)] # parent[u][i]=v : uの2^i世代上がv(なければ-1)
            # construct
            self.dfs(root)
            self.doubling()

        def dfs(self,u):
            for v in self.edges[u]:
                if self.depth[v]==float("inf"): # 訪れていない＝深さがINF(rootから連結な部分しか探索しない)
                    self.depth[v]=self.depth[u]+1
                    self.parent[v][0]=u
                    self.dfs(v)

        def doubling(self):
            for i in range(1,self.logn):
                for v in range(self.n):
                    if self.parent[v][i-1]!=-1: # 半分遡った時点で親が存在しなければ-1のままにしておく。-1が返るとlistの末尾の値の取得になってしまう
                        self.parent[v][i]=self.parent[self.parent[v][i-1]][i-1]

        def get(self,u,v): # uとvとのLCAを返す
            dd=self.depth[v]-self.depth[u]
            if dd<0: # vの方が深いようにする
                u,v=v,u
                dd=-dd

            for i in range(self.logn): # dd分だけvを遡らせる
                if dd&1: # 各bitで判定
                    v=self.parent[v][i]
                dd>>=1

            if u==v: return u # 高さ揃えた時点で一致してたら終わり

            for i in range(self.logn-1,-1,-1): # そうでなければ上から二分探索
                pu,pv=self.parent[u][i],self.parent[v][i]
                if pu!=pv:
                    u,v=pu,pv
            return self.parent[u][0]

    #%%
    # given data
    N=int(input())
    E=[set() for _ in range(N)]
    root=0
    for _ in range(N-1):
        u,v=map(lambda x:int(x)-1,input().split())
        E[u].add(v)
        E[v].add(u)
    lca=LCA(E,root)
    Q=int(input())
    for _ in range(Q):
        u,v=map(lambda x:int(x)-1,input().split())
        print(lca.depth[u]+lca.depth[v]-2*lca.depth[lca.get(u,v)]+1)
resolve()

Submission Info

Judge Result