Submission #1083102

#	Time	Username	Problem	Language	Result	Execution time	Memory
1083102		phong	Mergers (JOI19_mergers)	C++17	100 / 100	984 ms	216056 KiB

mergers

#include<bits/stdc++.h>

#define ll long long
#define int long long
const int nmax = 5e5 + 5, N = 1e5;
const ll oo = 1e18 + 1, base = 311;
const int lg = 19, M = 10;
const ll mod = 998244353, mod2 = 1e9 + 5277;
#define pii pair<int, int>
#define fi first
#define se second
#define endl "\n"
#define debug(a, n) for(int i = 1; i <= n; ++i) cout << a[i] << ' '; cout << "\n";
using namespace std;

int n, k, a[nmax];
vector<int> adj[nmax], one[nmax];
int st[nmax], en[nmax], tour[nmax], sc = 0;
int h[nmax], up[nmax][lg + 1], par[nmax];

void dfs_1(int u, int p){
    tour[++sc] = u;
    st[u] = sc;
    for(auto v : adj[u]){
        if(v == p) continue;
        h[v] = h[u] + 1;
        up[v][0] = u;
        for(int j = 1; j <= lg; ++j) up[v][j] = up[up[v][j - 1]][j - 1];
        dfs_1(v, u);
    }
    en[u] = sc;
}

int lca(int u, int v){
    if(h[u] != h[v]){
        if(h[u] < h[v]) swap(u, v);
        int k = h[u] - h[v];
        for(int j = 0; j <= lg; ++j){
            if(k >> j & 1){
                u = up[u][j];
            }
        }
    }
    if(u == v) return u;
    for(int j = __lg(h[u]); j >= 0; --j){
        if(up[u][j] != up[v][j]){
            u = up[u][j];
            v = up[v][j];
        }
    }
    return up[u][0];
}
int c[nmax];
void dfs_2(int u, int p){
    for(auto v : adj[u]){
        if(v == p) continue;
        dfs_2(v, u);
        c[u] += c[v];
    }
}
int r[nmax];

int get(int u){
    return r[u] ? r[u] = get(r[u]) : u;
}
void Union(int u, int v){
    u = get(u);
    v = get(v);
    if(u != v){
        r[u] = v;
    }
}
vector<int> gg[nmax], leaf;
int cnt = 0, deg = 0;
void dfs_3(int u, int p){
    st[u] = ++sc;
    deg++;
    bool ok = 1;
    for(auto v : gg[u]){
        if(v == p) continue;
        h[v] = h[u] + 1;
        up[v][0] = u;
        for(int j = 1; j <= lg; ++j) up[v][j] = up[up[v][j - 1]][j - 1];
        dfs_3(v, u);
        ok = 0;
    }
    if(ok) ++cnt, leaf.push_back(u);
}
main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
//    freopen("code.inp", "r", stdin);
//    freopen("code.out", "w", stdout);
    cin >> n >> k;
    for(int i = 1, u, v; i < n;++i){
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    for(int i =1; i <= n; ++i){
        cin >> a[i];
        one[a[i]].push_back(i);
    }
    dfs_1(1, -1);
    vector<int> tmp, two;
    for(int i = 1; i <= k; ++i){
        sort(one[i].begin(), one[i].end(), [](int a, int b){
             return st[a] < st[b];
             });
        tmp = one[i];
        for(int j = 1; j < one[i].size(); ++j){
            tmp.push_back(lca(one[i][j - 1], one[i][j]));
        }
        sort(tmp.begin(), tmp.end());
        tmp.erase(unique(tmp.begin(), tmp.end()), tmp.end());
        sort(tmp.begin(), tmp.end(), [](int a, int b){
             return st[a] < st[b];
            });

        for(int j = 0; j < tmp.size(); ++j){
            int u = tmp[j];
            if(j == 0) two.push_back(u);
            else{
                while(1){
                    int x = two.back();
                    if(st[x] <= st[u] && en[u] <= en[x])break;
                    two.pop_back();
                }
                int x = two.back();
                two.push_back(u);
                c[x]--;
                c[u]++;
            }
        }
        two.clear();
    }
    dfs_2(1, -1);
    for(int i = 2; i <= n; ++i){
        int x = up[i][0];
        if(c[i]){
            Union(i, x);
        }
    }
    for(int i = 1; i <= n; ++i){
        int u = get(i);
        for(auto v : adj[i]){
            int x = get(i), y = get(v);
            if(x != y){
                gg[x].push_back(y);
            }
        }
    }
    dfs_3(get(1), -1);
//    cout << deg<< ' ';
    if(deg == 1) cout << 0, exit(0);
    if(cnt & 1) cout << (cnt + 1) / 2;
    else{
        sort(leaf.begin(), leaf.end(),[](int a, int b){
             return st[a] < st[b];
             });
        int lc = lca(leaf[0], leaf.back());
        cnt /= 2;
        if(lc != get(1)) ++cnt;
        cout << cnt;
    }

}
/*
5 4
1 2
2 3
3 4
3 5
1 2 1 3 4


*/

Compilation message (stderr)

mergers.cpp:89:1: warning: ISO C++ forbids declaration of 'main' with no type [-Wreturn-type]
   89 | main(){
      | ^~~~
mergers.cpp: In function 'int main()':
mergers.cpp:111:26: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  111 |         for(int j = 1; j < one[i].size(); ++j){
      |                        ~~^~~~~~~~~~~~~~~
mergers.cpp:120:26: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  120 |         for(int j = 0; j < tmp.size(); ++j){
      |                        ~~^~~~~~~~~~~~
mergers.cpp:145:13: warning: unused variable 'u' [-Wunused-variable]
  145 |         int u = get(i);
      |             ^

• Subtask 1: 10 / 10
• Subtask 2: 24 / 24
• Subtask 3: 14 / 14
• Subtask 4: 22 / 22
• Subtask 5: 30 / 30