結果

問題 No.3348 Tree Balance
コンテスト
ユーザー ZOI-dayo
提出日時 2025-10-29 16:17:22
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
RE  
実行時間 -
コード長 3,818 bytes
コンパイル時間 356 ms
コンパイル使用メモリ 12,288 KB
実行使用メモリ 10,496 KB
最終ジャッジ日時 2025-11-13 21:03:33
合計ジャッジ時間 2,355 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other RE * 25
権限があれば一括ダウンロードができます

ソースコード

diff # 

import sys
# sortedcontainersライブラリをインポート
try:
    from sortedcontainers import SortedList
except ImportError:
    print("エラー: 'sortedcontainers' ライブラリが必要です。", file=sys.stderr)
    print("  pip install sortedcontainers  ", file=sys.stderr)
    sys.exit(1)

# 再帰深度の上限を増やす
sys.setrecursionlimit(200010) 

def main():
    # 高速I/O
    n_str = sys.stdin.readline()
    if not n_str:
        return
    n = int(n_str)
    
    w = list(map(int, sys.stdin.readline().split()))
    
    graph = [[] for _ in range(n)]
    for _ in range(n - 1):
        a, b = map(int, sys.stdin.readline().split())
        a -= 1
        b -= 1
        graph[a].append(b)
        graph[b].append(a)

    # --- dfs1: 各ノードを根とする部分木の重みの合計を計算 ---
    sum_subtree = list(w)
    
    def dfs1(i, p):
        for c in graph[i]:
            if c == p:
                continue
            sum_subtree[i] += dfs1(c, i)
        return sum_subtree[i]

    dfs1(0, -1)

    total = sum_subtree[0]
    ans = 1 << 60 

    # --- update: 答えを更新するヘルパー関数 ---
    def update(s1, s2):
        nonlocal ans
        s3 = total - s1 - s2
        if s3 < 0: # 3分割が不可能な場合は無視 (念のため)
            return
        mx = max(s1, s2, s3)
        mn = min(s1, s2, s3)
        ans = min(ans, mx - mn)

    # --- dfs2: Small-to-Large Merging (SortedListを使用) ---
    def dfs2(i, p):
        """
        SortedListを使い、O(N log^2 N)で最適解を探す。
        """
        # C++: new set<ll>() -> Python: SortedList()
        sums = SortedList()
        current = sum_subtree[i]
        remain = total - current

        for c in graph[i]:
            if c == p:
                continue
            
            # 子ノードの結果セット(SortedList)を取得
            res = dfs2(c, i)

            # --- Logic 1: (remain, *itr, current - *itr) ---
            # resは既にソート済みなので、ソート不要
            if res:
                # C++: res->lower_bound(current / 2)
                # O(log |res|)
                idx = res.bisect_left(current / 2)
                
                if idx < len(res):
                    update(remain, res[idx]) # C++の *itr
                if idx > 0:
                    update(remain, res[idx - 1]) # C++の *--itr

            # --- Small-to-Large Merging の最適化 ---
            # `sums` が常に `res` より大きい (か等しい) ようにスワップする
            if len(sums) < len(res):
                sums, res = res, sums

            # --- Logic 2: (e, *itr, total - e - *itr) ---
            # `sums` も既にソート済みなので、ソート不要
            if sums:
                # `res` (小さいセット) の各要素 e について
                # O(|res|)
                for e in res:
                    # O(log |sums|)
                    idx = sums.bisect_left((total - e) / 2)
                    
                    if idx < len(sums):
                        update(e, sums[idx])
                    if idx > 0:
                        update(e, sums[idx - 1])

            # --- マージ ---
            # C++: sums->insert(all(*res));
            # O(|res| * log |sums|)
            sums.update(res) 
            # Note: 
            # C++の set::insert(range) は O(|res| * log(|sums| + |res|))
            # SortedList.update() も O(|res| * log(|sums|))
            # よって、ここの計算量もC++とオーダーレベルで一致する

        # 最後に自分自身の部分木の合計もセットに追加
        # O(log |sums|)
        sums.add(current)
        return sums

    dfs2(0, -1)

    print(ans)

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