BOJ 22360 - UCPC 만들기

BOJ 22360 - UCPC 만들기

$H(u,c,p)$ 라는 해시를 관리한다.

각각 1,2,1 개 이상이 있다면 그만큼 모두 빼줘도 무관하다.

예를 들어, 2,5,4 개가 있다면 0,1,2 개에 대한 해시로 생각해도 무관하다.

어떤 centroid에서 자식들에 대한 해시를 업데이트 해주며, 이제 어떤 경로에서 $u,c,p$ 가 각각 있는것과 대응되는 해시값의 개수를 세주는 것이 정답이다.

int n;
vi ch;
int idx(char c) {
   if (c == 'U') return 0;
   if (c == 'C') return 1;
   if (c == 'P') return 2;
   assert(-1);
}
array<int, 3> normal(int u, int c, int p) {
   int q = min({u, p, c / 2});
   u -= q;
   c -= q * 2;
   p -= q;
   return {u, c, p};
}

inline ll H(int u, int c, int p) {
   auto nm = normal(u, c, p);
   u = nm[0], c = nm[1], p = nm[2];
   return 1ll * p + 1ll * 200001 * u + 1ll * 200001 * 200001 * c;
}
ll ans = 0;
struct centroid {
   int N;
   vi tree_size, vis;
   vvi _edges;
   centroid(int N) : N(N), tree_size(N), vis(N), _edges(N) {}
   void add_edge(int u, int v) { _edges[u].pb(v); }
   int get_size(int cur, int p) {
      tree_size[cur] = 1;
      for (int to: _edges[cur])if (to != p && !vis[to])tree_size[cur] += get_size(to, cur);
      return tree_size[cur];
   }
   int get_centroid(int cur, int p, int cap) {
      for (int to: _edges[cur]) if (to != p && !vis[to] && tree_size[to] * 2 > cap) return get_centroid(to, cur, cap);
      return cur;
   }

   // 반대쪽에 UUUUUUUCPC 가 있다고 하자
   // P 7개와 C 14개가 있어야 한다


   map<ll, ll> hashes, added;

   int cur_cent = -1;
   // centroid를 제외한 u,c,p, 의 개수만 관리
   void dfs(int cur, int u, int c, int p, int par) {
      auto nm = normal(u, c, p);
      u = nm[0], c = nm[1], p = nm[2];

      ll h = H(u + (ch[cur_cent] == 0), c + (ch[cur_cent] == 1), p + (ch[cur_cent] == 2));
      added[h]++;
      //debug("add", h, cur, u + (ch[cur_cent] == 0), c + (ch[cur_cent] == 1), p + (ch[cur_cent] == 2));

      int mx_q = max({u, p, (c + 1) / 2});
      int need_u = mx_q - u;
      int need_c = mx_q * 2 - c;
      int need_p = mx_q - p;

      ll target_h = H(need_u, need_c, need_p);
      //debug("query", need_u, need_c, need_p, target_h);
      if (hass(hashes, target_h)) {
         debug(cur, cur_cent, u, c, p, need_u, need_c, need_p);
         ans += hashes[target_h];
      }

      for (int to: _edges[cur]) {
         if (to != par && !vis[to]) {
            dfs(to, u + (ch[to] == 0), c + (ch[to] == 1), p + (ch[to] == 2), cur);
         }
      }
   }

   int build_tree(int cur, int p = -1) {
      cur = get_centroid(cur, -1, get_size(cur, -1));
      vis[cur] = 1;
      cur_cent = cur;

      hashes[H(ch[cur] == 0, ch[cur] == 1, ch[cur] == 2)]++;
      for (int to: _edges[cur]) {
         if (!vis[to]) {
            dfs(to, ch[to] == 0, ch[to] == 1, ch[to] == 2, cur);

            for (const auto &[h, cnt]: added) hashes[h] += cnt;
            added = {};
         }
      }
      hashes = {};

      for (int to: _edges[cur])if (!vis[to]) to = build_tree(to, cur);
      return cur;
   }
};
void solve() {
   cin >> n;
   string s;
   cin >> s;
   for (int i = 0; i < n; i++) ch.pb(idx(s[i]));
   centroid cent(n);
   for (int i = 0, u, v; i < n - 1; i++) {
      cin >> u >> v, u--, v--;
      cent.add_edge(u, v);
      cent.add_edge(v, u);
   }
   cent.build_tree(0);
   cout << ans;
}