Submission #559897

#	Time	Username	Problem	Language	Result	Execution time	Memory
559897		loctildore	Capital City (JOI20_capital_city)	C++14	0 / 100	506 ms	35592 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define f first
#define s second
//#define endl '\n'
#define all(x) begin(x), end(x)
int n,k,a,b,ts;
vector<int> grph[200069],towns[200069];
int col[200069],cnt[200069],curcnt[200069],done[200069];
int sz[200069],par[200069];
bool cent[200069];
void clear(int x, int lst=-1) {
  curcnt[col[x]]=0;
  done[col[x]]=0;
  for (auto i : grph[x]) {
    if (i==lst||cent[i]) {
      continue;
    }
    clear(i,x);
  }
}
int dfs(int x, int lst) {
  sz[x]=1;
  par[x]=lst;
  curcnt[col[x]]++;
  for (auto i : grph[x]) {
    if (i==lst||cent[i]) {
      continue;
    }
    sz[x]+=dfs(i,x);
  }
  return sz[x];
}
int fndcent(int x, int lst=-1) {
  for (auto i : grph[x]) {
    if (i==lst||cent[i]) {
      continue;
    }
    if (sz[i]*2>ts) {
      return fndcent(i,x);
    }
  }
  return x;
}
int solve(int x) {
  clear(x);
  ts=dfs(x,x);
  x=fndcent(x);
  //x is now the centroid
  stack<int> sk;
  sk.push(col[x]);
  int ans=-1;
  while (sk.size()) {
    int tmp=sk.top();
    sk.pop();
    if (done[tmp]) {
      continue;
    }
    done[tmp]=true;
    if (curcnt[tmp]!=cnt[tmp]) {
      ans=INT_MAX;
      break;
    }
    ans++;
    for (auto i:towns[tmp]) {
      sk.push(col[par[i]]);
    }
  }
  cent[x]=true;
  for (auto i:grph[x]) {
    if (!cent[i]) {
      ans=min(ans,solve(i));
    }
  }
  return ans;
}
int main()
{
  ios_base::sync_with_stdio(false);
  cin.tie(nullptr);
  cin>>n>>k;
  for (int i = 0; i < n-1; i++) {
    cin>>a>>b;
    a--;b--;
    grph[a].push_back(b);
    grph[b].push_back(a);
  }
  for (int i = 0; i < n; i++) {
    cin>>col[i];
    col[i]--;
    cnt[col[i]]++;
    towns[col[i]].push_back(i);
  }
  cout<<solve(0)<<endl;
  return 0;
}

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