Submission #627171

#	Time	Username	Problem	Language	Result	Execution time	Memory
627171		MohamedFaresNebili	Capital City (JOI20_capital_city)	C++14	100 / 100	553 ms	38332 KiB

capital_city

#include <bits/stdc++.h>

            using namespace std;

            using ll = long long;
            using ld = long double;
            using pi = pair<pair<int, int>, pair<int, int>>;

            #define ff first
            #define ss second
            #define pb push_back
            #define all(x) (x).begin(), (x).end()
            #define lb lower_bound

            int N, K, res = 1e9 + 7, C[200005];
            int cur[200005], par[200005], sz[200005];
            bool used[200005], seen[200005];
            vector<int> adj[200005], S[200005];

            int dfs_size(int v, int p) {
                sz[v] = 1; cur[C[v]] = 0; used[C[v]] = 0;
                for(auto u : adj[v]) {
                    if(u == p || seen[u]) continue;
                    sz[v] += dfs_size(u, v);
                }
                return sz[v];
            }
            int dfs_centroid(int v, int p, int _N) {
                for(auto u : adj[v]) {
                    if(u == p || seen[u]) continue;
                    if(sz[u] >= _N / 2) 
                        return dfs_centroid(u, v, _N);
                }
                return v;
            }
            void dfs(int v, int p) {
                par[v] = p; cur[C[v]]++; 
                for(auto u : adj[v]) {
                    if(u == p || seen[u]) continue;
                    dfs(u, v);
                }
            }
            void build(int v) {
                int _N = dfs_size(v, v);
                int centroid = dfs_centroid(v, v, _N);
                dfs(centroid, centroid); seen[centroid] = 1;
                int ans = 0;
                stack<int> st; st.push(C[centroid]);
                while(!st.empty()) {
                    int u = st.top(); st.pop();
                    if(used[u]) continue;
                    if(cur[u] != S[u].size()) { ans = 1e9 + 7; break; }
                    used[u] = 1; ans++;
                    for(auto e : S[u]) st.push(C[par[e]]);
                }
                res = min(res, ans - 1);
                for(auto u : adj[centroid]) {
                    if(seen[u]) continue;
                    build(u);
                }
            }

            int32_t main() {
                ios_base::sync_with_stdio(0); cin.tie(0); cout.tie(0);
                cin >> N >> K;
                for(int l = 0; l < N - 1; l++) {
                    int A, B; cin >> A >> B;
                    adj[A].push_back(B);
                    adj[B].push_back(A);
                }
                for(int l = 1; l <= N; l++)
                    cin >> C[l], S[C[l]].push_back(l);
                build(1); cout << res << "\n";
                return 0;
            }

Compilation message (stderr)

capital_city.cpp: In function 'void build(int)':
capital_city.cpp:52:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   52 |                     if(cur[u] != S[u].size()) { ans = 1e9 + 7; break; }
      |                        ~~~~~~~^~~~~~~~~~~~~~

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