BOJ 25638 - 트리와 경로 개수 쿼리

BOJ 25638 - 트리와 경로 개수 쿼리

Re-rooting Technic을 이용해서 어떤 정점 $u$가 루트가 됐을 때 각 자식 서브트리에서 $0, 1$의 개수를 관리하면 $O(E)$로 셀 수 있다.

void solve() {
   int n;
   cin >> n;
   vi a(n);
   fv(a);
   vi ans(n);
   vvi edges(n);
   for (int i = 0, u, v; i < n - 1; i++) {
      cin >> u >> v, u--, v--, edges[u].pb(v), edges[v].pb(u);
   }
   vvi dp(n, vi(2));
   function<void(int, int)> fn = [&](int i, int p) -> void {
      dp[i][a[i]]++;
      for (int j: edges[i]) {
         if (j == p) continue;
         fn(j, i);
         dp[i][0] += dp[j][0];
         dp[i][1] += dp[j][1];
      }
   };
   fn(0, -1);
   function<void(int, int)> fn2 = [&](int i, int p) -> void {
      int c0 = 0, c1 = 0;
      for (int j: edges[i]) {
         ans[i] += dp[j][0] * c1;
         ans[i] += dp[j][1] * c0;
         c0 += dp[j][0];
         c1 += dp[j][1];
      }

      for (int j: edges[i]) {
         if (j == p) continue;

         dp[i][0] -= dp[j][0];
         dp[i][1] -= dp[j][1];
         dp[j][0] += dp[i][0] /*- (a[i] == 0)*/;
         dp[j][1] += dp[i][1] /*- (a[i] == 1)*/;
         fn2(j, i);
         dp[j][1] -= dp[i][1] /*- (a[i] == 1)*/;
         dp[j][0] -= dp[i][0] /*- (a[i] == 0)*/;
         dp[i][0] += dp[j][0];
         dp[i][1] += dp[j][1];
      }
   };
   fn2(0, -1);
   int q;
   cin >> q;
   while (q--) {
      int u;
      cin >> u;
      cout << ans[u - 1] << endl;
   }
}