Submission #1159022

#	Time	Username	Problem	Language	Result	Execution time	Memory
1159022		fryingduc	Capital City (JOI20_capital_city)	C++20	0 / 100	302 ms	38800 KiB

capital_city

#include "bits/stdc++.h"

using namespace std;

#ifdef duc_debug
#include "bits/debug.h"
#else
#define debug(...)
#endif

const int maxn = 2e5 + 5;
int n, k;
vector<int> g[maxn];
vector<int> vc[maxn];
int c[maxn];
int color_sz[maxn];
vector<int> active;
int par[maxn], sz[maxn];
bool vis[maxn], del[maxn];
int res;

void re_dfs(int u, int prev) {
  active.push_back(u);
  sz[u] = 1;
  for (auto v:g[u]) {
    if (v == prev || del[v]) continue;
    re_dfs(v, u);
    sz[u] += sz[v];
  }
}

int find_centroid(int u, int prev, int n) {
  for (auto v:g[u]) {
    if (v == prev || del[v]) continue;
    if (sz[v] * 2 > n) {
      return find_centroid(v, u, n);
    }
  }
  return u;
}

int lab[maxn];

int find(int u) {
  return lab[u] < 0 ? u : lab[u] = find(lab[u]);
}

void join(int u, int v) {
  u = find(u), v = find(v);
  if (u == v) return;
  lab[u] += lab[v];
  lab[v] = u;
}

void dfs_cen(int u, int prev) {
  for (auto v:g[u]) {
    if (v == prev || del[v]) continue;
    par[v] = u;
    dfs_cen(v, u);
  }
}

void decompose(int u, int prev) {
  active.clear();
  re_dfs(u, prev);
  int cen = find_centroid(u, prev, sz[u]);
  for (auto x:active) {
    lab[x] = -1;
    vis[c[x]] = 0;
    vc[c[x]].push_back(x);
  }
  par[cen] = 0;
  dfs_cen(cen, 0);
  queue<int> q;
  vis[c[cen]] = 1;
  q.push(c[cen]);
  bool flag = 1;
  int tot = 0;
  while (!q.empty()) {
    int cur_color = q.front();
    q.pop();
    ++tot;
    if (color_sz[cur_color] != (int)vc[cur_color].size()) {
      flag = 0;
      break;
    }
    for (auto vertex:vc[cur_color]) {
      vertex = find(vertex);
      if (par[vertex]) {
        if (!vis[c[par[vertex]]]) {
          q.push(c[par[vertex]]);
          vis[c[par[vertex]]] = 1;
        }
        join(vertex, par[vertex]);
      }
    }
  }
  if (flag) {
    res = min(res, tot - 1);
  }
  del[cen] = 1;
  for (auto x:active) {
    vc[c[x]].clear();
  }
  for (auto v:g[cen]) {
    if (del[v] || v == prev) continue;
    decompose(v, cen);
  }
}

void solve() {
  cin >> n >> k;
  for (int i = 1; i < n; ++i) {
    int u, v; cin >> u >> v;
    g[u].push_back(v);
    g[v].push_back(u);
  }
  for (int i = 1; i <= n; ++i) {
    cin >> c[i];
    ++color_sz[c[i]];
  }
  res = n;
  decompose(1, 0);
  cout << res;  
}

signed main() {
  ios_base::sync_with_stdio(0);
  cin.tie(0);

  solve();

  return 0;
}

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