Submission #1205576

#	Time	Username	Problem	Language	Result	Execution time	Memory
1205576		repmann	Capital City (JOI20_capital_city)	C++20	100 / 100	543 ms	28736 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;
int N, K, glob, sol = INT_MAX;
int A[200096], D[200096], COUNT[200096];
bool V[200096];
int SIZE[200096], PARENT[200096];
vector <int> VG[200096], I[200096];
inline void SDFS(int node, int parent = 0)
{
  SIZE[node] = 1;
  for(int x : VG[node])
  {
    if(V[x] || (x == parent)) continue;
    SDFS(x, node);
    SIZE[node] += SIZE[x];
  }
  return;
}
inline int CDFS(int node, int target, int parent = 0)
{
  for(int x : VG[node])
  {
    if(V[x] || (SIZE[x] <= target) || (x == parent)) continue;
    return CDFS(x, target, node);
  }
  return node;
}
inline void DFS(int node, int parent = 0)
{
  if(D[A[node]] != glob)
  {
    I[A[node]] = vector <int>();
    D[A[node]] = glob;
  }
  I[A[node]].push_back(node);
  PARENT[node] = parent;
  for(int x : VG[node])
  {
    if(V[x] || (x == parent)) continue;
    DFS(x, node);
  }
  return;
}
inline void Centroid(int root)
{
  SDFS(root);
  int centroid = CDFS(root, SIZE[root] >> 1);
  V[centroid] = true;
//  cout << root << ' ' << centroid << '\n';
  glob++;
  DFS(centroid);
  glob++;
  queue <int> Q;
  int node = centroid, temp = 0;
  D[A[node]] = glob;
  if(COUNT[A[node]] == I[A[node]].size()) for(int x : I[A[node]]) Q.push(x);
  else temp = INT_MAX;
  while(Q.size())
  {
    node = Q.front();
    Q.pop();
    if(!PARENT[node]) continue;
    if(D[A[PARENT[node]]] != glob)
    {
      if(COUNT[A[PARENT[node]]] != I[A[PARENT[node]]].size()) {temp = INT_MAX; break;}
      temp++;
      D[A[PARENT[node]]] = glob;
      for(int x : I[A[PARENT[node]]]) Q.push(x);
    }
  }
  sol = min(temp, sol);
  for(int x : VG[centroid]) if(!V[x]) Centroid(x);
  return;
}
int main()
{
  ios_base::sync_with_stdio(0), cin.tie(0), cout.tie(0);
  cin >> N >> K;
  int u, v;
  for(int i = 2; i <= N; i++)
  {
    cin >> u >> v;
    VG[u].push_back(v);
    VG[v].push_back(u);
  }
  for(int i = 1; i <= N; i++)
  {
    cin >> A[i];
    COUNT[A[i]]++;
  }
  Centroid(1);
  cout << sol << '\n';

  return 0;
}

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