Submission #954992

#	Time	Username	Problem	Language	Result	Execution time	Memory
954992		abczz	Capital City (JOI20_capital_city)	C++14	100 / 100	622 ms	51300 KiB

capital_city

#include <iostream>
#include <vector>
#include <queue>
#define ll long long

using namespace std;

vector <ll> adj[200000];
bool C[200000];
vector <ll> V[200000];
ll n, k, a, b, cnt, f = 1e18, A[200000], B[200000], sz[200000], P[200000], visited[200000], G[200000];

void dfs_sz(ll u, ll p) {
  B[u] = cnt;
  sz[u] = 1;
  for (auto v : adj[u]) {
    if (v != p && !C[v]) {
      dfs_sz(v, u);
      sz[u] += sz[v];
    }
  }
}

ll find_cent(ll u, ll p, ll tot) {
  ll mx = -1, ret = -1;
  for (auto v : adj[u]) {
    if (v != p && !C[v]) {
      ret = max(ret, find_cent(v, u, tot));
      mx = max(mx, sz[v]);
    }
  }
  mx = max(mx, tot-sz[u]);
  if (mx <= tot/2) return u;
  else return ret;
}

void dfs_par(ll u, ll p) {
  for (auto v : adj[u]) {
    if (v != p && !C[v]) {
      dfs_par(v, u);
      P[v] = u;
    }
  }
}

void solve(ll u) {
  ll cur = 0;
  ++cnt;
  dfs_sz(u, -1);
  ll cent = find_cent(u, -1, sz[u]);
  C[cent] = 1;
  dfs_par(cent, -1);
  queue <ll> Q;
  bool ok = 1;
  G[A[cent]] = cnt;
  Q.push(A[cent]);
  while (!Q.empty()) {
    auto g = Q.front();
    Q.pop();
    ++cur;
    for (auto v : V[g]) {
      if (B[v] != cnt) {
        ok = 0;
        break;
      }
      ll x = v;
      while (visited[x] != cnt) {
        if (G[A[x]] != cnt) {
          G[A[x]] = cnt;
          Q.push(A[x]);
        }
        visited[x] = cnt;
        if (x == cent) break;
        x = P[x];
      }
    }
    if (!ok) break;
  }
  if (ok) f = min(f, cur);
  for (auto v : adj[cent]) {
    if (!C[v]) solve(v);
  }
}

int main() {
  cin >> n >> k;
  for (int i=0; i<n-1; ++i) {
    cin >> a >> b;
    --a, --b;
    adj[a].push_back(b);
    adj[b].push_back(a);
  }
  for (int i=0; i<n; ++i) {
    G[i] = -1;
    visited[i] = -1;
    cin >> A[i];
    --A[i];
    V[A[i]].push_back(i);
  }
  solve(0);
  cout << f-1 << '\n';
}

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