Submission #680839

#	Time	Username	Problem	Language	Result	Execution time	Memory
680839		tvladm2009	Capital City (JOI20_capital_city)	C++14	100 / 100	602 ms	42156 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;

const int N = (int) 2e5 + 7;
const int INF = (int) 1e9;
int c[N], timestamp[N], added[N], p[N], subarb[N];
vector<int> g[N], colors[N];
bool visited[N];

void getsize(int u, int par = 0)
{
  subarb[u] = 1;
  for (auto it : g[u])
  {
    if (visited[it] == false && it != par)
    {
      getsize(it, u);
      subarb[u] += subarb[it];
    }
  }
}

int getcentroid(int u, int total, int par = 0)
{
  for (auto it : g[u])
  {
    if (visited[it] == false && it != par && subarb[it] > total / 2)
    {
      return getcentroid(it, total, u);
    }
  }
  return u;
}

int Time = 0;

void dfs(int u, int par = 0)
{
  p[u] = par;
  timestamp[u] = Time;
  for (auto it : g[u])
  {
    if (visited[it] == false && it != par)
    {
      dfs(it, u);
    }
  }
}

int ret = INF;

void solve(int u)
{
  getsize(u);
  int centroid = getcentroid(u, subarb[u], 0);
  visited[centroid] = 1;
  Time++;
  timestamp[centroid] = Time;
  for (auto it : g[centroid])
  {
    if (visited[it] == false)
    {
      dfs(it, centroid);
    }
  }
  queue<int> q;
  bool ok = 1;
  q.push(c[centroid]);
  added[c[centroid]] = Time;
  int cnt = 1;
  while (!q.empty())
  {
    int color = q.front();
    q.pop();
    for (auto it : colors[color])
    {
      if (timestamp[it] != Time)
      {
        ok = 0;
        break;
      }
      if (it != centroid && added[c[p[it]]] != Time)
      {
        q.push(c[p[it]]);
        added[c[p[it]]] = Time;
        cnt++;
      }
    }
    if (ok == 0)
    {
      break;
    }
  }
  if (ok)
  {
    ret = min(ret, cnt - 1);
  }
  for (auto it : g[centroid])
  {
    if (visited[it] == false)
    {
      solve(it);
    }
  }
}

int main()
{
#ifdef ONPC
    freopen("input.txt", "r", stdin);
#endif // ONPC

#ifndef ONPC
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
#endif // ONPC

  int n, k;
  cin >> n >> k;
  for (int i = 1; i < n; i++)
  {
    int a, b;
    cin >> a >> b;
    g[a].push_back(b);
    g[b].push_back(a);
  }
  for (int i = 1; i <= n; i++)
  {
    cin >> c[i];
    colors[c[i]].push_back(i);
  }
  solve(1);
  cout << ret;
  return 0;
}

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