結果

問題 No.1833 Subway Planning
コンテスト
ユーザー neterukun
提出日時 2022-02-04 23:08:40
言語 PyPy3
(7.3.15)
結果
RE  
実行時間 -
コード長 3,635 bytes
記録
記録タグの例:
初AC ショートコード 純ショートコード 純主流ショートコード 最速実行時間
コンパイル時間 159 ms
コンパイル使用メモリ 82,432 KB
実行使用メモリ 237,444 KB
最終ジャッジ日時 2024-06-11 12:43:06
合計ジャッジ時間 20,163 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample RE * 3
other AC * 4 RE * 19
権限があれば一括ダウンロードができます

ソースコード

diff #
raw source code 

import sys
input = sys.stdin.buffer.readline


class UnionFind:
    def __init__(self, n):
        self.parent = [-1] * n
        self.n = n
        self.cnt = n

    def root(self, x):
        if self.parent[x] < 0:
            return x
        else:
            self.parent[x] = self.root(self.parent[x])
            return self.parent[x]

    def merge(self, x, y):
        x = self.root(x)
        y = self.root(y)
        if x == y:
            return False
        if self.parent[x] > self.parent[y]:
            x, y = y, x
        self.parent[x] += self.parent[y]
        self.parent[y] = x
        self.cnt -= 1
        return True

    def same(self, x, y):
        return self.root(x) == self.root(y)

    def size(self, x):
        return -self.parent[self.root(x)]

    def count(self):
        return self.cnt

    def groups(self):
        res = [[] for _ in range(self.n)]
        for i in range(self.n):
            res[self.root(i)].append(i)
        return [group for group in res if group]


def rerooting(n, edges, unit, merge, addnode):
    tree = [[] for i in range(n)]
    idxs = [[] for i in range(n)]
    for u, v in edges:
        idxs[u].append(len(tree[v]))
        idxs[v].append(len(tree[u]))
        tree[u].append(v)
        tree[v].append(u)
    sub = [[unit] * len(tree[v]) for v in range(n)]
    noderes = [unit] * n

    # topological sort
    tp_order = []
    par = [-1] * n
    for root in range(n):
        if par[root] != -1:
            continue
        stack = [root]
        while stack:
            v = stack.pop()
            tp_order.append(v)
            for nxt_v in tree[v]:
                if nxt_v == par[v]:
                    continue
                par[nxt_v] = v
                stack.append(nxt_v)

    # tree DP
    for v in reversed(tp_order[1:]):
        res = unit
        par_idx = -1
        for idx, nxt_v in enumerate(tree[v]):
            if nxt_v == par[v]:
                par_idx = idx
                continue
            res = merge(res, sub[v][idx])
        if par_idx != -1:
            sub[par[v]][idxs[v][par_idx]] = addnode(res, v)

    # rerooting DP
    for v in tp_order:
        acc_back = [unit] * len(tree[v])
        for i in reversed(range(1, len(acc_back))):
            acc_back[i - 1] = merge(sub[v][i], acc_back[i])
        acc_front = unit
        for idx, nxt_v in enumerate(tree[v]):
            res = addnode(merge(acc_front, acc_back[idx]), v)
            sub[nxt_v][idxs[v][idx]] = res
            acc_front = merge(acc_front, sub[v][idx])
        noderes[v] = addnode(acc_front, v)
    return sub
    return noderes


n = int(input())
edges = [list(map(int, input().split())) for _ in range(n - 1)]


es = []
tree = [[] for i in range(n)]
max_cost = 0
for u, v, cost in edges:
    u -= 1
    v -= 1
    max_cost = max(max_cost, cost)
    es.append((u, v))
    tree[u].append((v, cost))
    tree[v].append((u, cost))

vals = [0] * n 
for u, v, cost in edges:
    u -= 1
    v -= 1
    if cost == max_cost:
        vals[u] = 1
        vals[v] = 1

unit = 0
merge = lambda a, b: a + b
addnode = lambda val, v: val + vals[v]

sub = rerooting(n, es, unit, merge, addnode)
for v, res in enumerate(sub):
    if len(res) <= 1:
        continue
    res = sorted(res, reverse=True)
    if res[1] != 0:
        vals[v] = 1

is_path = True
ends = []
for v in range(n):
    if vals[v] == 0:
        continue
    cnt = 0
    for nxt_v, cost in tree[v]:
        if vals[nxt_v] >= 1:
            cnt += 1
    if cnt > 2:
        is_path = False
    if cnt == 1:
        ends.append(v)

if not is_path:
    print(max_cost)
    exit()

else:
    print(re)
0