RootedTree.py

class RootedTree:
    # G : adjcency list : 0 - indexed
    # N : the number of vertex
    def __init__(self,N,G):
        self.N = N
        self.G = G

        # for lca
        self.log_N = 20

        self.parents = [[ -1 for i in range(self.log_N)] for j in range(N)]
        self.depth = [-1] * self.N
        self.tree_initialize(0)

        for i in range(self.log_N-1):
            for j in range(N):
                if self.parents[j][i] == -1:
                    self.parents[j][i+1] = -1
                else:
                    self.parents[j][i+1] = self.parents[self.parents[j][i]][i]

    def tree_initialize(self, u):
        q = []
        q.append((u,-1,0))
        while q != []:
            u,parent,d = q.pop()

            self.depth[u] = d
            self.parents[u][0] = parent
            for v in self.G[u]:
                if v != parent:
                    q.append((v, u, d+1))

    def lca(self, u, v):
        if(self.depth[v] < self.depth[u]):
            u,v = v,u

        for i in range(self.log_N-1)[::-1]:
            if(((self.depth[v] - self.depth[u])>>i)&1):
                v = self.parents[v][i]

        if u==v:
            return u

        for i in range(self.log_N-1)[::-1]:
            if self.parents[u][i] != self.parents[v][i]:
                u = self.parents[u][i]
                v = self.parents[v][i]

        return self.parents[u][0]
    
    def dist(self, u, v):
        l = self.lca(u,v)
        return (self.depth[u] - self.depth[l]) + (self.depth[v] - self.depth[l])