Submission #602040

#	Time	Username	Problem	Language	Result	Execution time	Memory
602040		Soumya1	Capital City (JOI20_capital_city)	C++17	100 / 100	822 ms	43004 KiB

capital_city

#include <bits/stdc++.h>
#ifdef __LOCAL__
  #include <debug_local.h>
#endif
using namespace std;
const int mxN = 2e5 + 5;
vector<int> ad[mxN], l[mxN];
int city[mxN];
int ans = mxN;
int sz[mxN], par[mxN];
bool deleted[mxN], visc[mxN];
vector<int> vec;
void get_size(int u, int p = -1) {
  vec.push_back(u);
  visc[city[u]] = false;
  sz[u] = 1;
  for (int v : ad[u]) {
    if (v == p || deleted[v]) continue;
    get_size(v, u);
    par[v] = u;
    sz[u] += sz[v];
  }
}
int get_centroid(int u, int p, int n) {
  for (int v : ad[u]) {
    if (v == p || deleted[v]) continue;
    if (sz[v] > n / 2) return get_centroid(v, u, n); 
  }
  return u;
}
void solve(int s) {
  get_size(s);
  int c = get_centroid(s, -1, sz[s]);
  vec.clear();
  get_size(c);
  sort(vec.begin(), vec.end());
  int cur = 0;
  queue<int> q;
  visc[city[c]] = true;
  q.push(city[c]);
  while (!q.empty()) {
    cur++;
    int u = q.front();
    q.pop();
    for (int v : l[u]) {
      if (!binary_search(vec.begin(), vec.end(), v)) {
        cur = mxN;
        break;
      }
      else {
        if (v != c && !visc[city[par[v]]]) {
          q.push(city[par[v]]);
          visc[city[par[v]]] = true;
        }
      }
    }
  }
  ans = min(ans, cur);
  deleted[c] = true;
  for (int v : ad[c]) {
    if (!deleted[v]) solve(v);
  }
}
void testCase() {
  int n, k;
  cin >> n >> k;
  for (int i = 1; i < n; i++) {
    int u, v;
    cin >> u >> v;
    ad[u].push_back(v);
    ad[v].push_back(u);
  }
  for (int i = 1; i <= n; i++) {
    cin >> city[i];
    l[city[i]].push_back(i);
  }
  solve(1);
  cout << ans - 1 << "\n";
}
int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  testCase();
}

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