Submission #216891

#	Time	Username	Problem	Language	Result	Execution time	Memory
216891		_7_7_	Capital City (JOI20_capital_city)	C++14	100 / 100	801 ms	36012 KiB

capital_city

#include <bits/stdc++.h>                                           
 
//#define int long long
//#pragma GCC optimize("Ofast")
//#pragma comment(linker, "/stack:200000000")
//#pragma GCC target("sse,sse2,sse3,ssse3,sse4")
 
 
#define file(s) freopen(s".in","r",stdin); freopen(s".out","w",stdout);
#define forev(i, b, a) for(int i = (b); i >= (a); --i)
#define forn(i, a, b) for(int i = (a); i <= (b); ++i)
#define all(x) x.begin(), x.end()
#define sz(s) (int)s.size()
#define pb push_back
#define ppb pop_back
#define mp make_pair
#define s second
#define f first
 
 
using namespace std;
 
 
typedef pair < long long, long long > pll;    
typedef pair < int, int > pii;
typedef unsigned long long ull;         
typedef vector < pii > vpii;
typedef vector < int > vi;
typedef long double ldb;  
typedef long long ll;  
typedef double db;                             
 
 
const int dx[] = {1, -1, 0, 0}, dy[] = {0, 0, 1, -1}, block = 555;
const pii base = mp(1171, 3307), Mod = mp(1e9 + 7, 1e9 + 9);
const int inf = 1e9, maxn = 4e5 + 148, mod = 1e9 + 7, N = 2e5 + 53;
const db eps = 1e-12, pi = 3.14159265359;
const ll INF = 1e18;

int n, k, a, b, c[N], d[N], cnt[N], par[N], ans = inf, Sz;
bool was[N], ok, used[N], us[N], was1[N];
vi g[N], gg[N], cur;
queue < int > q;


void dfs (int v, int p = 0) {
    cur.pb(v);
    cnt[v] = 1;
    
    for (auto to : g[v])
        if (to != p && !was[to]) {
            dfs(to, v);
            cnt[v] += cnt[to];
        }
}

int Find (int v) {
    for (auto to : g[v])
        if (!was[to] && cnt[to] < cnt[v] && Sz - cnt[to] < cnt[to]) 
            return Find(to);
        
    return v;            
}

void dfs1 (int v, int p = 0) {
    par[v] = p;
    
    for (auto to : g[v])
        if (!was[to] && to != p)
            dfs1(to, v);
}


void dfs2 (int v) {
    was1[v] = 1;
    for (auto to : g[v])
        if (!was[to] && !was1[to])
            dfs2(to);
}

void Centroid (int v) {
    cur.clear();
    
    dfs(v);
    for (auto i : cur) {
        used[i] = us[c[i]] = 0;
        was1[i] = 0;
        gg[c[i]].clear();
    }
    
    for (auto i : cur)
        gg[c[i]].pb(i);
    
    Sz = cnt[v];
    v = Find(v);
    dfs1(v);
    
    q.push(c[v]);
    used[v] = 1;
    us[c[v]] = 1;
    
    int res = 0;
    ok = 1;
    while (sz(q)) {
        ++res;        
        int C = q.front();
        if (sz(gg[C]) != d[C]) 
            ok = 0;
        
        q.pop();
        
        for (auto i : gg[C]) {
            while (!used[i]) {
                used[i] = 1;
                if (!us[c[i]]) {
                    us[c[i]] = 1;
                    q.push(c[i]);
                }
                i = par[i];
            }
        }
    }
    
    if (ok) 
        ans = min(ans, res - 1);
    for (auto i : gg[c[v]])
        was[i] = 1;
    
    vi roots;
    for (auto i : cur)
        if (!was[i] && !was1[i]) {
            dfs2(i); 
            roots.pb(i);
        }
    for (auto i : roots)
        Centroid(i);
}


main () {
    scanf("%d%d", &n, &k);
    for (int i = 1; i < n; ++i) {
        scanf("%d%d", &a, &b);
        g[a].pb(b);
        g[b].pb(a);
    }
    
    for (int i = 1; i <= n; ++i) {
        scanf("%d", &c[i]); 
        ++d[c[i]];
    }
    
    Centroid(1);
    cout << ans << endl;
}

Compilation message (stderr)

capital_city.cpp:140:7: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
 main () {
       ^
capital_city.cpp: In function 'int main()':
capital_city.cpp:141:10: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
     scanf("%d%d", &n, &k);
     ~~~~~^~~~~~~~~~~~~~~~
capital_city.cpp:143:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
         scanf("%d%d", &a, &b);
         ~~~~~^~~~~~~~~~~~~~~~
capital_city.cpp:149:14: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
         scanf("%d", &c[i]); 
         ~~~~~^~~~~~~~~~~~~

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