image.png

BOJ 4315 - 나무 위의 구슬

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

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

ii 를 루트로 하는 서브트리에서 자식들을 CC 라고 하자. cc는 그 중 하나이다.

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

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

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

이 때의 비용도 동일하다. 넘기는 비용도 부족하거나 남는 개수만큼 똑같을 것이고 (1)=(2)(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;
}

Tags:

Categories:

Updated:

Comments