Submission #557879

#	Time	Username	Problem	Language	Result	Execution time	Memory
557879		600Mihnea	Capital City (JOI20_capital_city)	C++17	100 / 100	621 ms	39124 KiB

capital_city

#include <bits/stdc++.h>

bool home = 1;

using namespace std;

const int N=200000+7;
const int INF=(int)1e9+7;
int n,cntcols,col[N],sub[N],last_active[N],time_step,par[N],last_added[N],tt,print=INF;
vector<int> where_col[N];
vector<int> g[N];
bool blocked[N];
int total;

void build_sub(int a,int p=-1) {
  last_active[a]=time_step;
  sub[a]=1;
  for (auto &b:g[a]) {
    if(b==p||blocked[b]) continue;
    build_sub(b,a);
    sub[a]+=sub[b];
  }
}

int fcen(int a,int p=-1) {
  int mx=total-sub[a];
  for (auto &b:g[a]) {
    if(b==p||blocked[b]) continue;
    mx=max(mx,sub[b]);
  }
  if(mx<=total/2) return a;
  for (auto &b:g[a]) {
    if(b==p||blocked[b]) continue;
    if(sub[b]==mx) return fcen(b,a);
  }
  assert(0);
}

void dfs_par(int a,int p=-1) {
  par[a]=p;
  for (auto &b:g[a]) {
    if(b==p||blocked[b]) continue;
    dfs_par(b,a);
  }
}

void solve(int a) {
  time_step++;
  build_sub(a);
  total=sub[a];

  a=fcen(a);

  queue<int> colors;
  tt++;
  dfs_par(a);

  last_added[col[a]]=tt;
  colors.push(col[a]);
  bool valid=1;
  int sol=0;

  while (!colors.empty()&&valid) {
    int c=colors.front();
    colors.pop();
    for (auto &vertex:where_col[c]) {
      if (last_active[vertex]!=time_step) {
        valid=0;
        break;
      }
      if(vertex==a) continue;
      assert(1<=par[vertex]&&par[vertex]<=n);
      if(last_added[col[par[vertex]]]<tt){
        last_added[col[par[vertex]]]=tt;
        colors.push(col[par[vertex]]);
        sol++;
      }
    }
  }

  if (valid) {
    print=min(print,sol);
  }

  blocked[a]=1;
  for (auto &b:g[a]) {
    if(blocked[b]) continue;
    solve(b);
  }
}

signed main() {
#ifdef ONLINE_JUDGE
  home = 0;
#endif

  home=0;

  if (home) {
    freopen("I_am_iron_man", "r", stdin);
  }
  else {
    ios::sync_with_stdio(0); cin.tie(0);
  }


  cin>>n>>cntcols;
  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>>col[i];
    where_col[col[i]].push_back(i);
  }
  solve(1);
  cout<<print<<"\n";
}

Compilation message (stderr)

capital_city.cpp: In function 'int main()':
capital_city.cpp:100:12: warning: ignoring return value of 'FILE* freopen(const char*, const char*, FILE*)' declared with attribute 'warn_unused_result' [-Wunused-result]
  100 |     freopen("I_am_iron_man", "r", stdin);
      |     ~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~~~~~~

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