結果

問題 No.3348 Tree Balance
コンテスト
ユーザー Nzt3
提出日時 2025-11-13 12:04:24
言語 C++23
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 283 ms / 5,000 ms
コード長 2,741 bytes
コンパイル時間 3,201 ms
コンパイル使用メモリ 296,520 KB
実行使用メモリ 36,708 KB
最終ジャッジ日時 2025-11-13 21:28:35
合計ジャッジ時間 6,992 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 25
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>
using namespace std;

using ll = long long;
using vi = vector<int>;
using vl = vector<ll>;
#define rep3(i, a, b, c) for (ll i = (a); i < (b); i += (c))
#define rep2(i, a, b) rep3(i, a, b, 1)
#define rep1(i, n) rep2(i, 0, n)
#define rep0(n) rep1(aaaaa, n)
#define ov4(a, b, c, d, name, ...) name
#define rep(...) ov4(__VA_ARGS__, rep3, rep2, rep1, rep0)(__VA_ARGS__)

template <typename T, typename S> bool chmin(T &a, const S &b) {
  return a > b ? a = b, 1 : 0;
}

struct _ {
  _() { cin.tie(0)->sync_with_stdio(0), cout.tie(0); }
} __;

int main() {
  int N;
  cin >> N;
  vl W(N);
  for (ll &i : W)
    cin >> i;
  vector<vi> G(N);
  rep(N - 1) {
    int U, V;
    cin >> U >> V;
    --U, --V;
    G[U].push_back(V);
    G[V].push_back(U);
  }
  vl subtree = W;
  vi par(N);
  auto dfs_subtree = [&](auto f, int v, int p) -> ll {
    par[v] = p;
    for (int u : G[v]) {
      if (u != p)
        subtree[v] += f(f, u, v);
    }
    return subtree[v];
  };
  dfs_subtree(dfs_subtree, 0, -1);
  ll Wsum = subtree[0];
  auto dfs_centroid = [&](auto f, int v, int p) -> int {
    for (int u : G[v])
      if (u != p && subtree[u] * 2 >= Wsum)
        return f(f, u, v);
    return v;
  };
  int g = dfs_centroid(dfs_centroid, 0, -1);

  // 重心を根として部分木サイズを再計算
  subtree = W;
  dfs_subtree(dfs_subtree, g, -1);

  ll ans = subtree[g];
  {
    set<ll> cut1;
    // sを根とする部分木について調べる
    for (int s : G[g]) {
      queue<int> bfs;
      bfs.push(s);
      vl d;
      while (bfs.size()) {
        int v = bfs.front();
        bfs.pop();
        d.push_back(subtree[v]);
        for (int u : G[v]) {
          if (u != par[v]) {
            bfs.push(u);
          }
        }
      }

      // 重心を挟んで異なる部分木を切る場合
      // S1 - S2(G) - S3
      for (auto w : d) {
        ll mid = (Wsum - w + 1) / 2;
        auto it = cut1.lower_bound(mid);
        if (it != end(cut1)) {
          ll ma = max(*it, max(w, Wsum - w - (*it)));
          ll mi = min(*it, min(w, Wsum - w - (*it)));
          chmin(ans, ma - mi);
        }
        if (it != begin(cut1))
          it--;
        if (it != end(cut1)) {
          ll ma = max(*it, max(w, Wsum - w - (*it)));
          ll mi = min(*it, min(w, Wsum - w - (*it)));
          chmin(ans, ma - mi);
        }
      }

      // 重心から同じ部分木を2つに分ける場合
      // S1(G) - S2 - S3
      for (auto w : d) {
        ll ma = max(Wsum - subtree[s], max(w, subtree[s] - w));
        ll mi = min(Wsum - subtree[s], min(w, subtree[s] - w));
        chmin(ans, ma - mi);
      }
      for (auto w : d) {
        cut1.insert(w);
      }
    }
  }

  cout << ans << '\n';
}
0