BOJ 4315 - 나무 위의 구슬

BOJ 4315 - 나무 위의 구슬

$dp[i]$를 $i$를 루트로 하는 서브트리에서 $(1)$ 서브트리에 구슬의 총합이 노드의 수보다 많다면 모두 $1$개씩 나뉘가지고 루트에 나머지 것들을 몰아넣을 때의 최소 이동횟수, $(2)$ 그렇지 않다면 서브트리의 초기 구슬 상태를 유지한 채 부족한 만큼 $i$에 더 있다고 가정할 때 모두 $1$ 개 씩 있도록 분배할 때의 최소 이동횟수라고 정의하자.

$(1),~(2)$는 동치이다. $(2)$ 는 $(1)$ 을 수행한 것을 정확히 반대로 수행하면 되기 때문이다.

$i$ 를 루트로 하는 서브트리에서 자식들을 $C$ 라고 하자. $c$는 그 중 하나이다.

$c$ 를 루트로 하는 서브트리에서 구슬의 수가 노드의 수보다 많다면 결국 서브트리 $c$ 에 있는 남는 구슬들은 $c$ 로 이동된다음 그 수만큼 $i$ 로 이동되어야 한다.

이 때의 비용은 $dp[c] + \vert \text{서브트리 노드 개수}_c - \text{서브트리 구슬 개수}_c \vert$ 이다.

반대라면 $i$ 에 존재하는 구슬들을 $c$ 로 이동시켜줘야 한다. 여기서 $i$ 에 존재하는 구슬들은 원래 $i$ 에 있었거나 방금 구슬이 남는 자식들에서 넘어온 것이다.

이 때의 비용도 동일하다. 넘기는 비용도 부족하거나 남는 개수만큼 똑같을 것이고 $(1)=(2)$ 이기 때문이다.

따라서 이와 같이 트리 DP를 수행하면 된다.

void solve(int n) {
   vi a(n);
   vvi edges(n);
   for (int i = 0; i < n; i++) {
      int ni, d;
      cin >> ni, ni--;
      cin >> a[ni] >> d;
      for (int j = 0, c; j < d; j++) {
         cin >> c, c--;
         edges[ni].pb(c);
         edges[c].pb(ni);
      }
   }
   vi dp_root(n, -1), dp_sum(n), sub_cnt(n), diff(n);
   function<void(int, int)> fn = [&](int i, int p) -> void {
      dp_sum[i] = a[i];
      sub_cnt[i] = 1;
      for (int c: edges[i]) {
         if (c != p)
            fn(c, i), dp_sum[i] += dp_sum[c], sub_cnt[i] += sub_cnt[c];
      }
   };

   fn(0, -1);
   for (int i = 0; i < n; i++) diff[i] = abs(dp_sum[i] - sub_cnt[i]);

   function<int(int, int)> fn_root = [&](int i, int p) -> int {
      int &ret = dp_root[i];
      if (~ret) return ret;

      ret = 0;
      for (int c: edges[i]) {
         if (c != p)
            ret += diff[c] + fn_root(c, i);
      }
      return ret;
   };
   int ans = fn_root(0, -1);

   cout << ans << endl;
}

signed main() {
   fastio

   while (1) {
      int n;
      cin >> n;
      if (!n) break;
      solve(n);
   }
   return 0;
}