Submission #1196704

#	Time	Username	Problem	Language	Result	Execution time	Memory
1196704		rxlfd314	Capital City (JOI20_capital_city)	C++20	100 / 100	563 ms	30752 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_szs = [&](auto &&self, const int i, const int p) -> void {
    szs[i] = 1;
    nodes.push_back(i);
    for (const auto &j : adj[i]) {
      if (j == p || dead[j])
        continue;
      self(self, j, i);
      szs[i] += szs[j];
    }
  };
  auto centroid = [&](int i, const int n) {
    int p = -1;
    bool flag = true;
    while (flag) {
      flag = false;
      for (const auto &j : adj[i])
        if (j != p && !dead[j] && 2 * szs[j] > n) {
          p = i;
          i = j;
          flag = true;
          break;
        }
    }
    return i;
  };
  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);
  vt<int> seen(N);
  DSU uf(K);
  auto dfs = [&](auto &&self, const int root) -> void {
    dfs_szs(dfs_szs, root, -1);
    const int c = centroid(root, szs[root]);
    dfs_par(dfs_par, c, -1);
    #ifdef DEBUG
    cout << "centroid: " << c + 1 << '\n';
    #endif
    for (const auto &i : nodes)
      ind_list[C[i]].push_back(i);
    
    queue<int> qu;
    for (const auto &i : ind_list[C[c]])
      qu.push(i), seen[i] = 1;
    bool good = 1;
    int res = 0;
    while (size(qu) && good) {
      const int i = qu.front();
      qu.pop();
      good &= size(ind_list[C[i]]) == cnt[C[i]];
      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);
        }
      }
    }
    if (good)
      ans = min(ans, res);
    for (const auto &i : nodes) {
      uf.reset(C[i]);
      ind_list[C[i]].clear();
    }
    for (const auto &i : nodes)
      seen[i] = 0;
    nodes.clear();
    dead[c] = 1;
    for (const auto &j : adj[c])
      if (!dead[j])
        self(self, j);
  };
  
  dfs(dfs, 0);
  cout << ans << '\n';
}

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