Submission #228597

#	Time	Username	Problem	Language	Result	Execution time	Memory
228597		DS007	Mergers (JOI19_mergers)	C++14	70 / 100	343 ms	262148 KiB

mergers

#include <bits/stdc++.h>
using namespace std;
#define int long long

const int N = 5e5 + 1;

int dsu[N], s[N], first[N], dep[N], par[N], last[N];
vector<int> adj[N], adj_[N], cities[N], euler, depth;
pair<int, int> e[N];
int sparse[N * 2][22];
int n, k, c = 0;

// DSU begins
int find(int v) {
    if (dsu[v] == v)
        return v;
    else
        return dsu[v] = find(dsu[v]);
}

void un(int u, int v) {
    u = find(u), v = find(v);
    if (u == v)
        return;
    else if (dep[u] < dep[v])
        dsu[v] = u;
    else
        assert(dep[u] != dep[v]), dsu[u] = v;
}
// DSU ends

// LCA begins
void pre(int v, int p = -1, int d = 0) {
    first[v] = c++;
    dep[v] = d;
    par[v] = p;
    euler.push_back(v);
    depth.push_back(d);

    for (int i : adj[v]) {
        if (i != p) {
            pre(i, v, d + 1);
            euler.push_back(v);
            depth.push_back(d);
            c++;
        }
    }

    last[v] = c - 1;
}

int merge(int a, int b) {
    return depth[a] < depth[b] ? a : b;
}

void build_sparse() {
    for (int i = 0; i < euler.size(); i++)
        sparse[i][0] = i;

    for (int i = 1; i <= log2(euler.size()); i++) {
        for (int j = 0; j < euler.size() && j - 1 + (1 << i) < euler.size(); j++)
            sparse[j][i] = merge(sparse[j][i - 1], sparse[j + (1 << (i - 1))][i - 1]);
    }
}

int lca(int u, int v) {
    if (u == v) return u;
    int f1 = min(first[u], first[v]), f2 = max(first[v], first[u]), diff = f2 - f1;
    int dx = log2(diff);
    return euler[merge(sparse[f1][dx], sparse[f2 - (1 << dx)][dx])];
}
// LCA ends

void unite(int low, int high) {
    low = find(low), high = find(high);
    while (dep[low] > dep[high])
        un(low, high), low = find(par[low]);
}

void solveTestCase() {
    cin >> n >> k;

    for (int i = 0; i < n - 1; i++) {
        int u, v;
        cin >> u >> v;
        u--, v--;
        e[i] = {u, v};
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    euler.reserve(N * 2);
    depth.reserve(N * 2);
    iota(dsu, dsu + N, 0);
    pre(0);
    build_sparse();

    for (int i = 0; i < n; i++) {
        cin >> s[i];
        cities[s[i]].push_back(i);
    }

    for (int i = 0; i <= k; i++) {
        if (cities[i].size() <= 1)
            continue;

        int high = lca(cities[i][0], cities[i][1]);
        for (int j = 1; j < cities[i].size(); j++)
            high = lca(high, cities[i][j]);

        for (int j : cities[i])
            unite(j, high);
    }

    for (int i = 0; i < n - 1; i++) {
        if (find(e[i].first) != find(e[i].second))
            adj_[find(e[i].first)].push_back(find(e[i].second)), adj_[find(e[i].second)].push_back(find(e[i].first));
    }

    int leaves = 0;
    for (int i = 0; i < n; i++) {
        if (adj_[i].size() == 1)
            leaves++;
    }

    cout << (leaves + 1) / 2;
}

signed main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int test = 1;
    // cin >> test;
    for (int i = 1; i <= test; i++)
        solveTestCase();
}

Compilation message (stderr)

mergers.cpp: In function 'void build_sparse()':
mergers.cpp:57:23: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
     for (int i = 0; i < euler.size(); i++)
                     ~~^~~~~~~~~~~~~~
mergers.cpp:61:27: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
         for (int j = 0; j < euler.size() && j - 1 + (1 << i) < euler.size(); j++)
                         ~~^~~~~~~~~~~~~~
mergers.cpp:61:62: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
         for (int j = 0; j < euler.size() && j - 1 + (1 << i) < euler.size(); j++)
                                             ~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~
mergers.cpp: In function 'void solveTestCase()':
mergers.cpp:108:27: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
         for (int j = 1; j < cities[i].size(); j++)
                         ~~^~~~~~~~~~~~~~~~~~

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