結果

問題 No.1094 木登り / Climbing tree
ユーザー lloyz
提出日時 2022-10-14 00:27:56
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 1,526 ms / 2,000 ms
コード長 3,257 bytes
記録
記録タグの例:
初AC ショートコード 純ショートコード 純主流ショートコード 最速実行時間
コンパイル時間 231 ms
コンパイル使用メモリ 82,100 KB
実行使用メモリ 133,236 KB
最終ジャッジ日時 2024-06-26 12:16:39
合計ジャッジ時間 32,996 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1
other AC * 26
権限があれば一括ダウンロードができます

ソースコード

diff #
raw source code 

class HeavyLightDecomposition:
    # HL 分解 (HLD)
    def __init__(self, E, root=1):
        # E は双方向に辺を張った木で、破壊的に有向グラフに変更される
        # O(N)
        self.E = E
        self.root = root
        self.N = N = len(E)  # 頂点数
        self.Parent = [-1] * N  # 頂点番号 v -> 親ノード
        self.Size = [-1] * N  # 頂点番号 v -> 部分木のサイズ
        self.dfs1()

        self.Mapping = [-1] * N  # 頂点番号 v -> 内部インデックス
        self.Head = list(range(N))  # 頂点番号 v -> v を含む heavy path の左端の頂点番号
        self.Depth = [0] * N  # 頂点番号 v -> 深さ（root から v までの距離）
        self.dfs2()

    def dfs1(self):
        E = self.E
        Parent, Size = self.Parent, self.Size
        Path = [self.root]
        Idx_edge = [0]
        while Path:
            v = Path[-1]
            idx_edge = Idx_edge[-1]
            Ev = E[v]
            if idx_edge != len(Ev):
                # 行きがけ・通りがけ 辺の数だけ実行される
                u, c = Ev[idx_edge]
                Idx_edge[-1] += 1
                E[u].remove((v, c))  # 有向グラフならここをコメントアウトする
                Parent[u] = v
                Path.append(u)
                Idx_edge.append(0)
            else:
                # 帰りがけ 頂点の数だけ実行される
                if len(Ev) >= 2:
                    ma = -1
                    argmax = None
                    for i, (u, _) in enumerate(Ev):
                        if Size[u] > ma:
                            ma = Size[u]
                            argmax = i
                    Ev[0], Ev[argmax] = Ev[argmax], Ev[0]
                Size[v] = sum(Size[u] for u, _ in Ev) + 1
                Path.pop()
                Idx_edge.pop()
    
    def dfs2(self):
        E = self.E
        Mapping = self.Mapping
        Head = self.Head
        Depth = self.Depth
        k = 0
        St = [self.root]
        while St:
            v = St.pop()
            Mapping[v] = k
            k += 1
            Ev = E[v]
            if Ev:
                Head[Ev[0][0]] = Head[v]
                St += [u for u, _ in Ev[::-1]]
                for u, c in Ev:
                    Depth[u] = Depth[v] + c  # distance を使わないのならここをコメントアウトする
    
    def lca(self, v, u):
        # O(logN)
        Parent = self.Parent
        Mapping = self.Mapping
        Head = self.Head
        while True:
            if Mapping[v] > Mapping[u]:
                v, u = u, v  # v の方を根に近くする
            if Head[v] == Head[u]:
                return v
            u = Parent[Head[u]]
    
    def distance(self, v, u):
        # O(logN)
        Depth = self.Depth
        return Depth[v] + Depth[u] - 2 * Depth[self.lca(v, u)]

from collections import deque

n = int(input())
E = [[] for _ in range(n + 1)]
EC = [[] for _ in range(n + 1)]
for _ in range(n - 1):
    a, b, c = map(int, input().split())
    E[a].append((b, c))
    E[b].append((a, c))
HL = HeavyLightDecomposition(E)

q = int(input())
for _ in range(q):
    s, t = map(int, input().split())
    print(HL.distance(s, t))
0