Submission #235711

#	Time	Username	Problem	Language	Result	Execution time	Memory
235711		shashwatchandra	Capital City (JOI20_capital_city)	C++17	100 / 100	1202 ms	55272 KiB

capital_city

/*input
12 4
7 9
1 3
4 6
2 4
10 12
1 2
2 10
11 1
2 8
5 3
6 7
3
1
1
2
4
3
3
2
2
3
4
4
*/
#include<bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;

#define int long long 
#define double long double
#define f first
#define s second
#define mp make_pair
#define pb push_back

#define RE(i,n) for (int i = 1; i <= n; i++)
#define RED(i,n) for (int i = n; i > 0; i--)
#define REPS(i,n) for(int i = 1; (i*i) <= n; i++)
#define REP(i,n) for (int i = 0; i < (int)n; i++)
#define FOR(i,a,b) for (int i = a; i < b; i++)
#define REPD(i,n) for (int i = n-1; i >= 0; i--)
#define FORD(i,a,b) for (int i = a; i >= b; i--)

#define all(v) v.begin(),v.end()
#define pii pair<int,int>
#define vi vector<int>
#define vvi vector<vi>
#define print(arr) for (auto it = arr.begin(); it != arr.end(); ++it) cout << *it << " "; cout << endl;
#define debug(x) cout << x << endl;
#define debug2(x,y) cout << x << " " << y << endl;
#define debug3(x,y,z) cout << x << " " << y << " " << z << endl;

typedef tree<
int,
null_type,
less<int>,
rb_tree_tag,
tree_order_statistics_node_update>
ordered_set;

const int INF = 1e18+1;
const int MOD = 1e9+7;
const double PI = 3.14159265358979323846264338;

int raise(int a,int n,int m = MOD){
  if(n == 0)return 1;
  if(n == 1)return a;
  int x = 1;
    x *= raise(a,n/2,m);
    x %= m;
    x *= x;
    x %= m;
    if(n%2)x*= a;
    x %= m;
    return x;
}

int floor1(int n,int k){
    if(n%k == 0 || n >= 0)return n/k;
    return (n/k)-1;
}

int ceil1(int n,int k){
    return floor1(n+k-1,k);
}

const int N = 2e5+1;
int n;
int k;
int col[N];
vector<int> all[N];
bool broken[N];
int freq[N];
int par[N];
set<int> adj[N];
int nn = 0;
int sub[N];
bool thisdone[N];
bool fine = 1;
int ans = INF;
int cur;
queue<int> q;

void dfs1(int u,int p){
    nn++;
    freq[col[u]]++;
    sub[u] = 1;
    for(int v:adj[u]){
      if(v == p)continue;
      dfs1(v,u);
      sub[u] += sub[v];
    }
}

void closecheck(int u,int p){
    if(freq[col[u]] != all[col[u]].size())broken[col[u]] = 1;
    par[u] = p;
    for(int v:adj[u]){
      if(v == p)continue;
      closecheck(v,u);
    }
}

void dellnow(int u,int p){
    freq[col[u]]--;
    thisdone[col[u]] = 0;
    for(int v:adj[u]){
      if(v == p)continue;
      dellnow(v,u);
    }
}

int findcen(int u,int p){
    for(int v:adj[u]){
      if(v == p)continue;
      if(sub[v] > nn/2)return findcen(v,u);
    }
    return u;
}


void addcolor(int c){
    if(broken[c]){fine = 0;return;}
    if(thisdone[c])return;
    thisdone[c] = 1;
    cur++;
    for(int v:all[c]){
      q.push(v);
    }
}

void cendfs(int u,int p = -1){
    nn = 0;
    dfs1(u,-1);
    int centroid = findcen(u,-1);
    // check here
    closecheck(centroid,-1);
    //cout << centroid << endl;
    fine = 1;
    cur = 0;
    addcolor(col[centroid]);
    while(!q.empty()){
      int x= q.front();
      q.pop();
      if(par[x] != -1)addcolor(col[par[x]]);
    }
    //cout << fine << " " << cur << endl;
    if(fine)ans = min(ans,cur);
    // 
    dellnow(centroid,-1);
    for(int x:adj[centroid]){
      adj[x].erase(centroid);
      cendfs(x,centroid);
    }
}

void solve(){
    cin >> n >> k;
    RE(i,n-1){
      int a,b;cin >> a >> b;
      adj[a].insert(b);
      adj[b].insert(a);
    }
    RE(i,n){
      cin >> col[i];
      all[col[i]].pb(i);
    }
    cendfs(1,-1);
    cout << ans-1;
}

signed main(){
  ios_base::sync_with_stdio(0);cin.tie(0);cout.tie(0);
  //freopen(".in","r",stdin);freopen(".out","w",stdout);
  int t = 1;
  //cin >> t;
  while(t--){
    solve();
  }
  return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'void closecheck(long long int, long long int)':
capital_city.cpp:120:21: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     if(freq[col[u]] != all[col[u]].size())broken[col[u]] = 1;
        ~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~

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