Submission #225216

#	Time	Username	Problem	Language	Result	Execution time	Memory
225216		quocnguyen1012	Capital City (JOI20_capital_city)	C++14	100 / 100	982 ms	40936 KiB

capital_city

#include <bits/stdc++.h>

#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define eb emplace_back
#define ar array

using namespace std;
typedef long long ll;
typedef pair<int, int> ii;

const int maxn = 2e5 + 5, inf = 1e9;

int N, K;
vector<int> same[maxn], adj[maxn];
int sz[maxn], del[maxn], vis[maxn], cvis[maxn], col[maxn], par[maxn];
vector<int> vec;
int ret = inf, num[maxn];

void prepare(int u, int p)
{
  vec.pb(col[u]);
  par[u] = p;
  same[col[u]].pb(u);
  vis[u] = cvis[col[u]] = 0;
  for(int v : adj[u]) if(v != p && !del[v]){
    prepare(v, u);
  }
}

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

int findcen(int u, int p, int siz)
{
  for(int v : adj[u]) if(v != p && !del[v]){
    if(sz[v] > siz / 2) return findcen(v, u, siz);
  }
  return u;
}

void decomp(int u)
{
  findsz(u, u);
  u = findcen(u, u, sz[u]);
  prepare(u, u);
  del[u] = true;

  queue<int> q;
  q.push(col[u]);
  cvis[col[u]] = true;
  int cnt = 0;
  while(q.size()){
    int c = q.front(); q.pop();
    ++cnt;
    for(int v : same[c]){
      while(!vis[v]){
        vis[v] = true;
        if(!cvis[col[v]]){
          cvis[col[v]] = true;
          q.push(col[v]);
        }
        v = par[v];
      }
    }
  }
  bool flag = true;
  for(int c : vec){
    if(cvis[c] && num[c] != same[c].size())
      flag = false;
  }
  //cerr << u << ' ' << cnt << '\n';
  if(flag)
    ret = min(ret, cnt - 1);
  for(int c : vec)
    same[c].clear();
  vec.clear();
  for(int v : adj[u])
    if(!del[v])
      decomp(v);
}

signed main(void)
{
  ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
  #ifdef LOCAL
    freopen("A.INP", "r", stdin);
    freopen("A.OUT", "w", stdout);
  #endif // LOCAL
  cin >> N >> K;
  for(int i = 1; i < N; ++i){
    int u, v; cin >> u >> v;
    adj[u].eb(v); adj[v].eb(u);
  }
  for(int i = 1; i <= N; ++i){
    cin >> col[i];
    ++num[col[i]];
  }
  decomp(1);
  cout << ret;
}

Compilation message (stderr)

capital_city.cpp: In function 'void decomp(int)':
capital_city.cpp:77:26: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     if(cvis[c] && num[c] != same[c].size())
                   ~~~~~~~^~~~~~~~~~~~~~~~~

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