BOJ 26972 - Barn Tree

BOJ 26972 - Barn Tree

부모에서부터 자식으로 가는 간선이 부모 $\to$ 자식에게 결과적으로 값을 줘야한다면 비용으로 $+c$ 를 갖고, 반대라면 $-c$ 를 갖는다고 하자.

어떤 간선에서 $u \rightarrow v$ 로 줄 수 있다면 $u$ 를 큐에 넣고 BFS를 돌린다.

이제 $h_v$가 늘어났으므로 $v$를 큐에넣고 BFS에 돌린다.

이렇게하면 시간안에 올바른 순서를 얻을 수 있다.

BFS에 진행하며 큐에 들어가는 개수의 상한은 $O(E)$이고 항상 자신이 자식에게 줄 수 있을 때 주는 순서대로 답을 구성할 수 있기 때문이다.

struct edge {
   int to, rev_idx, cost;
};

void solve() {
   int n;
   cin >> n;
   vi h(n), par(n);
   fv(h);
   ll tot = acc(h), solo = tot / n;
   vector<vector<edge>> edges(n);
   for (int i = 0, u, v; i < n - 1; i++) {
      cin >> u >> v, u--, v--;
      edges[u].pb({v, sz(edges[v])});
      edges[v].pb({u, sz(edges[u]) - 1});
   }
   vi sub_sum(n), sub_size(n), edge_cost(n, 0);
   function<void(int, int)> fn = [&](int i, int p) -> void {
      par[i] = p;
      sub_sum[i] += h[i];
      sub_size[i]++;
      for (edge &e: edges[i]) {
         if (e.to ^ p) {
            fn(e.to, i);
            sub_sum[i] += sub_sum[e.to];
            sub_size[i] += sub_size[e.to];
            // p -> child
            int cost = sub_size[e.to] * solo - sub_sum[e.to];
            e.cost = cost;
            edges[e.to][e.rev_idx].cost = -cost;
         }
      }
   };
   fn(0, -1);
   queue<int> q;
   for (int i = 0; i < n; i++) {
      for (edge &e: edges[i]) {
         if (h[i] >= e.cost) {
            q.push(i);
            break;
         }
      }
   }
   vector<array<int, 3>> ans;
   while (sz(q)) {
      int cur = q.front();
      q.pop();
      for (edge &e: edges[cur]) {
         if (e.cost > 0 && e.cost <= h[cur]) {
            h[e.to] += e.cost;
            h[cur] -= e.cost;
            ans.pb({cur + 1, e.to + 1, e.cost});
            e.cost = 0;
            edges[e.to][e.rev_idx].cost = 0;
            q.push(e.to);
         }
      }
   }
   cout << sz(ans) << endl;
   for (auto &[s, e, c]: ans) cout << s << ' ' << e << ' ' << c << endl;
}