Submission #216829

#	Time	Username	Problem	Language	Result	Execution time	Memory
216829		IOrtroiii	Capital City (JOI20_capital_city)	C++14	100 / 100	691 ms	36072 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;

const int MAXN = 200200;

int N, K;
vector<int> adj[MAXN];
int city[MAXN];
int par[MAXN];
vector<int> towns[MAXN];
int sz[MAXN];
bool delt[MAXN];
vector<int> curComp;
int ans = MAXN;

void dfsSz(int v, int p) {
   sz[v] = 1;
   for (int u : adj[v]) if (u != p && !delt[u]) {
      dfsSz(u, v);
      sz[v] += sz[u];
   }
}

int findCentr(int v, int p, int r) {
   for (int u : adj[v]) if (u != p && !delt[u] && 2 * sz[u] >= sz[r]) {
      return findCentr(u, v, r);
   }
   return v;
}

int dfsPar(int v, int p) {
   par[v] = p;
   curComp.emplace_back(v);
   for (int u : adj[v]) if (u != p && !delt[u]) {
      dfsPar(u, v);
   }
}

int root[MAXN];

int getRoot(int v) {
   if (root[v] != v) {
      root[v] = getRoot(root[v]);
   }
   return root[v];
}

bool merge(int v, int u) {
   v = getRoot(v);
   u = getRoot(u);
   if (v == u) {
      return false;
   } else {
      root[u] = v;
      return true;
   }
}

int cnt[MAXN];

void solve(int v) {
   dfsSz(v, v);
   v = findCentr(v, v, v);
   dfsPar(v, v);
   for (int u : curComp) {
      ++cnt[city[u]];
   }
   bool good = cnt[city[v]] == int(towns[city[v]].size());
   int cans = 0;
   queue<int> q;
   if (good) {
      for (int u : towns[city[v]]) {
      q.emplace(u);
      }
      while (!q.empty() && good) {
         int u = q.front(); q.pop();
         if (merge(city[u], city[par[u]])) {
            good &= (cnt[city[par[u]]] == int(towns[city[par[u]]].size()));
            if (!good) break;
            cans++;
            for (int z : towns[city[par[u]]]) {
               q.emplace(z);
            }
         }
      }
   }
   if (good) ans = min(ans, cans);
   for (int z : curComp) {
      root[city[z]] = city[z];
      cnt[city[z]] = 0;
   }
   curComp = {};
   delt[v] = true;
   for (int u : adj[v]) if (!delt[u]) solve(u);
}

int main() {
   ios_base::sync_with_stdio(false); cin.tie(nullptr);
   cin >> N >> K;
   for (int i = 1; i < N; ++i) {
      int v, u;
      cin >> v >> u;
      adj[v].emplace_back(u);
      adj[u].emplace_back(v);
   }
   for (int v = 1; v <= N; ++v) {
      cin >> city[v];
      towns[city[v]].emplace_back(v);
   }
   iota(root + 1, root + K + 1, 1);
   solve(1);
   cout << ans << "\n";
   return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'int dfsPar(int, int)':
capital_city.cpp:38:1: warning: no return statement in function returning non-void [-Wreturn-type]
 }
 ^

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