結果

問題 No.1094 木登り / Climbing tree
ユーザー kept1994kept1994
提出日時 2022-11-23 14:43:15
言語 PyPy3
(7.3.15)
結果
AC  
実行時間 1,541 ms / 2,000 ms
コード長 3,400 bytes
コンパイル時間 492 ms
コンパイル使用メモリ 81,616 KB
実行使用メモリ 141,304 KB
最終ジャッジ日時 2023-10-24 23:59:56
合計ジャッジ時間 39,478 ms
ジャッジサーバーID
(参考情報)
judge14 / judge12
このコードへのチャレンジ(β)

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 41 ms
52,368 KB
testcase_01 AC 1,498 ms
130,456 KB
testcase_02 AC 293 ms
141,304 KB
testcase_03 AC 584 ms
80,956 KB
testcase_04 AC 545 ms
100,336 KB
testcase_05 AC 668 ms
123,212 KB
testcase_06 AC 959 ms
94,396 KB
testcase_07 AC 1,541 ms
130,804 KB
testcase_08 AC 1,489 ms
131,808 KB
testcase_09 AC 1,503 ms
129,392 KB
testcase_10 AC 1,473 ms
129,796 KB
testcase_11 AC 1,472 ms
130,448 KB
testcase_12 AC 1,458 ms
130,164 KB
testcase_13 AC 1,450 ms
129,836 KB
testcase_14 AC 1,452 ms
129,372 KB
testcase_15 AC 912 ms
90,848 KB
testcase_16 AC 1,072 ms
126,304 KB
testcase_17 AC 991 ms
105,628 KB
testcase_18 AC 913 ms
98,276 KB
testcase_19 AC 1,008 ms
117,152 KB
testcase_20 AC 1,539 ms
129,400 KB
testcase_21 AC 1,044 ms
108,096 KB
testcase_22 AC 1,523 ms
129,760 KB
testcase_23 AC 1,490 ms
130,844 KB
testcase_24 AC 1,469 ms
129,520 KB
testcase_25 AC 1,496 ms
130,660 KB
testcase_26 AC 1,493 ms
130,640 KB
権限があれば一括ダウンロードができます

ソースコード

diff #

#!/usr/bin/env python3
import sys

class LcaDoubling:
    # 木であれば任意の点を根と見做せる。
    def __init__(self, N, root=0):
        self.N = N
        self.root = root
        self.G = [[] for _ in range(N)]
        self.depths = [-1] * N
        self.distances = [-1] * N
        self.ancestors = []
        return
    
    def addEdge(self, fromNode: int, toNode: int, cost: int):
        self.G[fromNode].append((cost, toNode))
        # print("Really directed Graph?")
        return
    
    def build(self):
        prevAncestors = self._dfs()
        self.ancestors.append(prevAncestors)
        d = 1
        max_depth = max(self.depths)
        while d < max_depth:
            nextAncestors = [prevAncestors[p] for p in prevAncestors]
            self.ancestors.append(nextAncestors)
            d <<= 1
            prevAncestors = nextAncestors
        return

    def _dfs(self):
        q = [(self.root, -1, 0, 0)]
        directAncestors = [-1] * (self.N + 1)  # 頂点数より1個長くし、存在しないことを-1で表す。末尾(-1)要素は常に-1
        while q:
            now, parent, dep, dist = q.pop()
            directAncestors[now] = parent
            self.depths[now] = dep
            self.distances[now] = dist
            for cost, next in self.G[now]:
                if next != parent:
                    q.append((next, now, dep + 1, dist + cost))
        return directAncestors
 
    def getLca(self, nodeA: int, nodeB: int):
        depthA, depthB = self.depths[nodeA], self.depths[nodeB]
        if depthA > depthB:
            nodeA, nodeB = nodeB, nodeA
            depthA, depthB = depthB, depthA
        
        # 2ノードを同じ深さまで揃える。
        tu = nodeA
        tv = self.upstream(nodeB, depthB - depthA)

        # 遡上させて行き2つが衝突する位置が共通祖先。
        if nodeA == tv:
            return nodeA
        for k in range(depthA.bit_length() - 1, -1, -1):
            mu = self.ancestors[k][tu]
            mv = self.ancestors[k][tv]
            if mu != mv:
                tu = mu
                tv = mv
        lca = self.ancestors[0][tu]
        assert lca == self.ancestors[0][tv]
        return lca
 
    # 2つのノードの間の距離を返す。
    def getDistance(self, nodeA, nodeB):
        lca = self.getLca(nodeA, nodeB)
        return self.distances[nodeA] + self.distances[nodeB] - 2 * self.distances[lca]

    # targetNodeが2つのノード間のパス上に存在するかを返す。
    def isOnPath(self, nodeA: int, nodeB: int, evalNode: int):
        return self.getDistance(nodeA, nodeB) == self.getDistance(nodeA, evalNode) + self.getDistance(evalNode, nodeB) 

    # ノードvからk個遡上したノードを返す。
    def upstream(self, v, k):
        i = 0
        while k:
            if k & 1:
                v = self.ancestors[i][v]
            k >>= 1
            i += 1
        return v

def main():
    N = int(input())
    ld = LcaDoubling(N)
    for _ in range(N - 1):
        a, b, c = map(int, input().split())
        ld.addEdge(fromNode=a - 1, toNode=b - 1, cost=c)
        ld.addEdge(fromNode=b - 1, toNode=a - 1, cost=c)
    ld.build()
    Q = int(input())
    for _ in range(Q):
        s, t = map(int, input().split())
        print(ld.getDistance(nodeA=s - 1, nodeB=t - 1))

if __name__ == '__main__':
    main()
0