Submission #1196700

#	Time	Username	Problem	Language	Result	Execution time	Memory
1196700		rxlfd314	Capital City (JOI20_capital_city)	C++20	11 / 100	3094 ms	27432 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using ari2 = array<int, 2>;
using ari3 = array<int, 3>;
using arl2 = array<ll, 2>;
using arl3 = array<ll, 3>;
template <class T> using vt = vector<T>;
#define all(x) begin(x), end(x)
#define size(x) (int((x).size()))
#define REP(a,b,c,d) for(int a=(b);(d)>0?a<=(c):a>=(c);a+=(d))
#define FOR(a,b,c) REP(a,b,c,1)
#define ROF(a,b,c) REP(a,b,c,-1)

struct DSU {
  vt<int> uf;
  DSU(const int n) : uf(n, -1) {}
  int find(const int i) {
    return uf[i] < 0 ? i : uf[i] = find(uf[i]);
  }
  bool unite(int a, int b) {
    if ((a = find(a)) == (b = find(b)))
      return false;
    if (uf[a] > uf[b])
      swap(a, b);
    uf[a] += uf[b];
    uf[b] = a;
    return true;
  }
  void reset(const int i) {
    uf[i] = -1;
  }
};

signed main() {
#ifndef DEBUG
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
#endif
  int N, K;
  cin >> N >> K;
  vt<vt<int>> adj(N);
  FOR(_, 2, N) {
    int a, b;
    cin >> a >> b, a--, b--;
    adj[a].push_back(b);
    adj[b].push_back(a);
  }
  vt<int> C(N), cnt(N);
  FOR(i, 0, N-1) {
    cin >> C[i], C[i]--;
    cnt[C[i]]++;
  }
  
  vt<int> dead(N), szs(N), nodes, parent(N);
  int ans = N;
  auto dfs_par = [&](auto &&self, const int i, const int p) -> void {
    parent[i] = p;
    for (const auto &j : adj[i])
      if (j != p && !dead[j])
        self(self, j, i);
  };
  vt<vt<int>> ind_list(K);
  FOR(i, 0, N-1)
    ind_list[C[i]].push_back(i);
  vt<int> seen(N);
  DSU uf(K);
  auto dfs = [&](auto &&self, const int root) -> void {
    dfs_par(dfs_par, root, -1);
    queue<int> qu;
    for (const auto &i : ind_list[C[root]])
      qu.push(i), seen[i] = 1;
    int res = 0;
    while (size(qu)) {
      const int i = qu.front();
      qu.pop();
      if (parent[i] >= 0 && !seen[parent[i]] && uf.unite(C[i], C[parent[i]])) {
        res++;
        for (const auto &j : ind_list[C[parent[i]]]) {
          seen[j] = 1;
          qu.push(j);
        }
      }
    }
    ans = min(ans, res);
    FOR(i, 0, K-1)
      uf.reset(i);
    FOR(i, 0, N-1)
      seen[i] = 0;
  };
  FOR(i, 0, N-1)
    dfs(dfs, i);
  cout << ans << '\n';
}

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 0 / 30
• Subtask 4: 0 / 59