結果

問題 No.386 貪欲な領主
ユーザー rpy3cpprpy3cpp
提出日時 2016-07-01 23:32:03
言語 PyPy2
(7.3.15)
結果
AC  
実行時間 1,431 ms / 2,000 ms
コード長 3,109 bytes
コンパイル時間 303 ms
コンパイル使用メモリ 77,728 KB
実行使用メモリ 358,748 KB
最終ジャッジ日時 2023-08-02 20:56:48
合計ジャッジ時間 9,439 ms
ジャッジサーバーID
（参考情報） 		judge13 / judge14
このコードへのチャレンジ（β）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 72 ms
76,272 KB
testcase_01 AC 68 ms
76,544 KB
testcase_02 AC 70 ms
76,352 KB
testcase_03 AC 73 ms
76,388 KB
testcase_04 AC 1,369 ms
358,748 KB
testcase_05 AC 1,405 ms
192,888 KB
testcase_06 AC 1,431 ms
198,152 KB
testcase_07 AC 116 ms
81,432 KB
testcase_08 AC 324 ms
105,660 KB
testcase_09 AC 131 ms
82,256 KB
testcase_10 AC 70 ms
76,716 KB
testcase_11 AC 71 ms
76,604 KB
testcase_12 AC 108 ms
80,564 KB
testcase_13 AC 211 ms
86,180 KB
testcase_14 AC 898 ms
183,196 KB
testcase_15 AC 1,142 ms
318,684 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys

sys.setrecursionlimit(10**6)

input = raw_input
range = xrange


def read_data():
    N = int(input())
    Es = [[] for i in range(N)]
    for i in range(N - 1):
        a, b = map(int, input().split())
        Es[a].append(b)
        Es[b].append(a)
    Us = [int(input()) for i in range(N)]
    M = int(input())
    moves = [list(map(int, input().split())) for m in range(M)]
    return N, Es, Us, M, moves


class DisjointSet():
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n
    
    def union(self, x, y):
        self._link(self.find_set(x), self.find_set(y))
    
    def _link(self, x, y):
        if self.rank[x] > self.rank[y]:
            self.parent[y] = x
        else:
            self.parent[x] = y
            if self.rank[x] == self.rank[y]:
                self.rank[y] += 1
    
    def find_set(self, x):
        xp = self.parent[x]
        if xp != x:
            self.parent[x] = self.find_set(xp)
        return self.parent[x]

class Tree():
    def __init__(self, N, Es, root, Us):
        self.n = N
        self.root = root
        self.child = [[] for i in range(N)]
        self.cum_cost = [0 for i in range(N)]
        self._set_child(Es, Us)
    
    def _set_child(self, Es, Us):
        que = [self.root]
        visited = [False] * self.n
        self.cum_cost[self.root] = Us[self.root]
        while que:
            v = que.pop()
            cum_cost_v = self.cum_cost[v]
            for u in Es[v]:
                if visited[u]:
                    continue
                self.child[v].append(u)
                self.cum_cost[u] = cum_cost_v + Us[u]
                que.append(u)
            visited[v] = True        

class LCATarjan():
    def __init__(self, tree):
        self.n = tree.n
        self.root = tree.root
        self.child = tree.child
        self.ancestor = list(range(self.n))
        self.visited = [False] * self.n
        self.ds = DisjointSet(self.n)
        self.lca = dict()
    
    def find(self, pairs):
        self._preprocess(pairs)
        self._LCA(self.root)
        return self.lca

    def _preprocess(self, pairs):
        self.pairs = [[] for node in range(self.n)]
        for u, v in pairs:
            self.pairs[u].append(v)
            self.pairs[v].append(u)
    
    def _LCA(self, u):
        self.ancestor[self.ds.find_set(u)] = u
        for v in self.child[u]:
            self._LCA(v)
            self.ds.union(u, v)
            self.ancestor[self.ds.find_set(u)] = u
        self.visited[u] = True
        for v in self.pairs[u]:
            if self.visited[v]:
                self.lca[u, v] = self.ancestor[self.ds.find_set(v)]
                self.lca[v, u] = self.lca[u, v]

    

def solve(N, Es, Us, M, moves):
    tree = Tree(N, Es, 0, Us)
    cum_cost = tree.cum_cost
    pairs = [(a, b) for a, b, c in moves]
    lcat = LCATarjan(tree)
    lca = lcat.find(pairs)
    tax = 0
    for a, b, c in moves:
        v = lca[a, b]
        tax += (cum_cost[a] + cum_cost[b] - cum_cost[v] * 2 + Us[v]) * c
    return tax

pars = read_data()
print(solve(*pars))
0