BOJ 28296 - 물류창고

BOJ 28296 - 물류창고

간선이 가중치가 큰 것부터 보며 그 간선이 이어짐으로써 연결되는 동일한 회사들이 있다면

그 회사의 개수가 한 쪽에 $L$, 다른쪽에 $R$개라면 $LRw$ 를 그 회사의 정답에 더해줄 수 있다.

Small to Large Technic으로 map 을 관리해주며 DSU에서 합쳐주면 된다.

ll ans[50000];

struct DSU {
   vector<int> p;
   vector<map<int, int>> cnt;
   DSU(int n) : p(n, -1), cnt(n) {}
   int gp(int n) {
      if (p[n] < 0) return n;
      return p[n] = gp(p[n]);
   }
   void merge(int a, int b, int w) {
      a = gp(a), b = gp(b);
      if (a == b) return;
      if (size(a) > size(b)) swap(a, b);
      p[b] += p[a];
      p[a] = b;

      // a -> b
      if (sz(cnt[a]) > sz(cnt[b])) swap(cnt[a], cnt[b]);
      for (const auto &[team, c]: cnt[a]) {
         if (hass(cnt[b], team)) {
            ans[team] += 1ll * w * cnt[b][team] * c;
         }
         cnt[b][team] += c;
      }
   }
   bool is_merged(int a, int b) { return gp(a) == gp(b); }
   int size(int n) { return -p[gp(n)]; }
};
void solve() {
   int n, k, m;
   cin >> n >> k >> m;
   vi c(n);
   fv(c);
   for (int &i: c) i--;
   vector<array<int, 3>> edges;
   while (m--) {
      int x, y, w;
      cin >> x >> y >> w, x--, y--;
      edges.pb({w, x, y,});
   }
   DSU dsu(n);
   for (int i = 0; i < n; i++) dsu.cnt[i][c[i]] = 1;
   sort(all(edges), greater<>());
   for (const auto &[w, a, b]: edges) {
      dsu.merge(a, b, w);
   }
   for (int i = 0; i < k; i++) cout << ans[i] << endl;
}