Submission #898225

#	Time	Username	Problem	Language	Result	Execution time	Memory
898225		NeroZein	Capital City (JOI20_capital_city)	C++17	100 / 100	609 ms	42256 KiB

capital_city

#include "bits/stdc++.h"
using namespace std;

#ifdef Nero
#include "Deb.h"
#else
#define deb(...)
#endif

const int N = 2e5 + 5;

int pr[N];
int sz[N];
int tot[N];
int color[N];
bool viscol[N];
bool blocked[N];
vector<int> g[N];
vector<int> ofcolor[N];

void find_sizes(int v, int p, int flag) {
  sz[v] = 1;
  pr[v] = p;
  if (flag == 1) {
    ofcolor[color[v]].push_back(v);     
  } else if (flag == -1) {
    ofcolor[color[v]].pop_back();
  }
  for (int u : g[v]) {
    if (!blocked[u] && u != p) {
      find_sizes(u, v, flag);
      sz[v] += sz[u];
    }
  }
}

int find_centroid(int v, int p, int x) {
  for (int u : g[v]) {
    if (!blocked[u] && u != p && sz[u] > x / 2) {
      return find_centroid(u, v, x); 
    }
  }
  return v;
}

int decompose(int root) {
  find_sizes(root, root, 0); 
  int centroid = find_centroid(root, root, sz[root]); 
  find_sizes(centroid, centroid, 1); 
  queue<int> que;
  viscol[color[centroid]] = true; 
  vector<int> used = {color[centroid]}; 
  for (int i : ofcolor[color[centroid]]) {
    que.push(i); 
  }
  while (!que.empty()) {
    int v = que.front(); 
    que.pop();
    if (!viscol[color[pr[v]]]) { 
      used.push_back(color[pr[v]]); 
      viscol[color[pr[v]]] = true;
      for (int i : ofcolor[color[pr[v]]]) {
        que.push(i); 
      }
    }
  }
  bool ok = true; 
  for (int i : used) {  
    viscol[i] = false; 
    ok &= ((int) ofcolor[i].size() == tot[i]);
  }
  find_sizes(centroid, centroid, -1); 
  int ret = (ok ? (int) used.size() : N);
  blocked[centroid] = true; 
  for (int u : g[centroid]) {
    if (!blocked[u]) {
      ret = min(ret, decompose(u)); 
    }
  }
  return ret; 
}

int main() {
  ios::sync_with_stdio(false);
  cin.tie(nullptr);
  int n, k;
  cin >> n >> k;
  for(int i = 0; i < n - 1; ++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 >> color[i];
    tot[color[i]]++; 
  }
  int ans = decompose(1);
  cout << ans - 1 << '\n';
  return 0;
}

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