結果

問題 No.3348 Tree Balance
コンテスト
ユーザー ZOI-dayo
提出日時 2025-10-25 16:24:45
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
TLE  
実行時間 -
コード長 4,530 bytes
コンパイル時間 88 ms
コンパイル使用メモリ 12,416 KB
実行使用メモリ 129,548 KB
最終ジャッジ日時 2025-11-13 20:53:31
合計ジャッジ時間 12,831 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other TLE * 1 -- * 24
権限があれば一括ダウンロードができます

ソースコード

diff # 

# By Gemini

import sys

def solve():
    # 再帰深度を深く設定
    # N が大きい場合 (例: 2*10^5) に備える
    sys.setrecursionlimit(200010)

    # --- 入力処理 ---
    try:
        N = int(sys.stdin.readline())
        W = list(map(int, sys.stdin.readline().split()))
        
        adj = [[] for _ in range(N)]
        edges = [] # 辺 (u, v) を (0-indexed) で保存
        for _ in range(N - 1):
            A, B = map(int, sys.stdin.readline().split())
            u, v = A - 1, B - 1
            adj[u].append(v)
            adj[v].append(u)
            edges.append((u, v))

    except EOFError:
        return
    except Exception as e:
        # print(f"Input Error: {e}", file=sys.stderr)
        return

    # N < 3 の場合、辺を2本選べないが、問題文の前提 (N-1本の木) から
    # N >= 3 と想定される。
    if N < 3:
        print(0)
        return

    # --- 全頂点の重みの総和 ---
    total_sum = sum(W)

    # --- 根付き木 (頂点0を根) の情報を計算 (DFS) ---
    parent = [-1] * N       # parent[v]: 頂点vの親
    subtree_sum = [0] * N   # subtree_sum[v]: vを根とする部分木の重み総和
    in_time = [-1] * N      # in_time[v]: DFSの行きがけ順
    out_time = [-1] * N     # out_time[v]: DFSの帰りがけ順
    timer = [0]             # DFSの時刻 (ミュータブルなリストとして渡す)

    def dfs(v, p):
        parent[v] = p
        in_time[v] = timer[0]
        timer[0] += 1
        
        s = W[v] # 自身の重み
        
        for u in adj[v]:
            if u == p:
                continue
            dfs(u, v)
            s += subtree_sum[u] # 子の部分木の重みを加算
            
        subtree_sum[v] = s
        out_time[v] = timer[0]
        timer[0] += 1

    # 頂点0を根としてDFSを実行
    dfs(0, -1)

    # --- 祖先判定関数 (O(1)) ---
    # u が v の祖先 (u == v も含む) かどうか
    def is_ancestor(u, v):
        return in_time[u] <= in_time[v] and out_time[v] <= out_time[u]

    # --- 各辺 (N-1本) に対応する「子ノード」を特定 ---
    # 辺 (u, p) を削除したとき、p側 (根側) と u側 (子側) に分かれる。
    # u (子ノード) の subtree_sum[u] が切り出される部分木の重み和となる。
    child_nodes = []
    for u, v in edges:
        if parent[v] == u:
            # v が u の子
            child_nodes.append(v)
        else:
            # u が v の子
            child_nodes.append(u)
    
    # child_nodes には N-1 個の頂点インデックスが入る

    min_diff = float('inf')

    # --- 2本の辺 (i, j) を全探索 O(N^2) ---
    num_edges = N - 1
    for i in range(num_edges):
        v1 = child_nodes[i]
        S_A = subtree_sum[v1] # 辺iの削除で切り出される部分木の重み和

        for j in range(i + 1, num_edges):
            v2 = child_nodes[j]
            S_B = subtree_sum[v2] # 辺jの削除で切り出される部分木の重み和

            sums = []
            
            # 場合分け:
            # 1. v1 が v2 の祖先 (辺iが辺jより根に近い)
            if is_ancestor(v1, v2):
                S1 = S_B                     # v2の部分木
                S2 = S_A - S_B               # v1の部分木からv2の部分木を除いたもの
                S3 = total_sum - S_A         # v1の親側 (全体からv1の部分木を除いたもの)
                sums = [S1, S2, S3]
            
            # 2. v2 が v1 の祖先 (辺jが辺iより根に近い)
            elif is_ancestor(v2, v1):
                S1 = S_A                     # v1の部分木
                S2 = S_B - S_A               # v2の部分木からv1の部分木を除いたもの
                S3 = total_sum - S_B         # v2の親側 (全体からv2の部分木を除いたもの)
                sums = [S1, S2, S3]
            
            # 3. v1 と v2 が祖先-子孫関係にない (並列な部分木)
            else:
                S1 = S_A                     # v1の部分木
                S2 = S_B                     # v2の部分木
                S3 = total_sum - S_A - S_B   # それ以外 (根を含む部分)
                sums = [S1, S2, S3]

            # (最大値 - 最小値) の最小値を更新
            diff = max(sums) - min(sums)
            if diff < min_diff:
                min_diff = diff

    # --- 出力 ---
    print(min_diff)

if __name__ == "__main__":
    solve()
0