Submission #711836

#	Time	Username	Problem	Language	Result	Execution time	Memory
711836		Jarif_Rahman	Capital City (JOI20_capital_city)	C++17	100 / 100	533 ms	81116 KiB

capital_city

#include <bits/stdc++.h>
#define pb push_back
#define f first
#define sc second
using namespace std;
typedef long long int ll;
typedef string str;

int n, k;
vector<vector<int>> tree;
vector<pair<int, int>> edges;
vector<int> C, cnt;
vector<map<int, int>> cities;
vector<set<int>> graph;
vector<vector<int>> graph_r;

vector<int> order;
vector<bool> bl;
vector<int> comp_id;
vector<vector<int>> comp;

void dfs(int nd, int ss){
    for(int x: tree[nd]) if(x != ss) dfs(x, nd);

    for(int x: tree[nd]) if(x != ss && C[x] != C[nd] && cities[x][C[nd]] > 0)
        graph[C[nd]].insert(C[x]);
    for(int x: tree[nd]) if(x != ss){
        if(cities[x].size() > cities[nd].size()) swap(cities[x], cities[nd]);
        for(auto [c, y]: cities[x]) cities[nd][c]+=y;
        cities[x].clear();
    }
    cities[nd][C[nd]]++;
    if(ss != -1 && C[ss] != C[nd] && cities[nd][C[nd]] < cnt[C[nd]]) 
        graph[C[nd]].insert(C[ss]);
}

void dfs2(int nd){
    if(bl[nd]) return;
    bl[nd] = 1;
    for(int x: graph[nd]) dfs2(x);
    order.pb(nd);
}
void dfs3(int nd){
    if(bl[nd]) return;
    bl[nd] = 1;
    comp_id[nd] = int(comp.size())-1;
    comp.back().pb(nd);
    for(int x: graph_r[nd]) dfs3(x);
}

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);

    cin >> n >> k;

    tree.resize(n);
    C.resize(n);
    cnt.resize(k);
    cities.resize(n);
    graph.resize(k);
    graph_r.resize(k);

    for(int i = 0; i < n-1; i++){
        int a, b; cin >> a >> b; a--, b--;
        tree[a].pb(b);
        tree[b].pb(a);
        edges.pb({a, b});
    }
    for(int &x: C) cin >> x, x--, cnt[x]++;

    dfs(0, -1);

    for(int i = 0; i < k; i++){
        for(int x: graph[i]) graph_r[x].pb(i);
    }

    bl.assign(k, 0);
    comp_id.resize(k);

    for(int i = 0; i < k; i++) if(!bl[i]) dfs2(i);
    reverse(order.begin(), order.end());

    bl.assign(k, 0);
    for(int x: order) if(!bl[x]){
        comp.pb({});
        dfs3(x);
    }

    int ans = n-1;
    for(int i = 0; i < comp.size(); i++){
        bool ok = 1;
        for(int x: comp[i]) for(int y: graph[x]) if(comp_id[y] != i) ok = 0;
        if(ok) ans = min(ans, int(comp[i].size())-1);
    }

    cout << ans << "\n";
}

Compilation message (stderr)

capital_city.cpp: In function 'int main()':
capital_city.cpp:91:22: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   91 |     for(int i = 0; i < comp.size(); i++){
      |                    ~~^~~~~~~~~~~~~

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