Submission #706050

#	Time	Username	Problem	Language	Result	Execution time	Memory
706050		pakhomovee	Capital City (JOI20_capital_city)	C++17	11 / 100	3032 ms	524288 KiB

capital_city

#include <iostream>
#include <vector>
#include <algorithm>
#include <cmath>
#include <set>
#include <stdlib.h>
#include <map>
#include <random>
#include <cstring>
#include <functional>
#include <iomanip>
#include <cassert>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <array>

using namespace std;

const int inf = 1e9;
const int maxn = 1 << 18;

int par[maxn];
int tin[maxn];
int tout[maxn];
int up[20][maxn];
int timer = 0;
void dfs(int v, vector<vector<int>> &g, int p) {
    tin[v] = timer++;
    up[0][v] = p;
    for (int i = 1; i < 20; ++i) {
        up[i][v] = up[i - 1][up[i - 1][v]];
    }
    for (int u : g[v]) {
        if (u != p) {
            par[u] = v;
            dfs(u, g, v);
        }
    }
    tout[v] = timer++;
}

bool isa(int u, int v) {
    return tin[u] <= tin[v] && tout[v] <= tout[u];
}

int lca(int u, int v) {
    if (isa(u, v)) return u;
    if (isa(v, u)) return v;
    for (int i = 19; i >= 0; --i) {
        if (!isa(u, v)) u = up[i][u];
    }
    return up[0][u];
}

vector<int> ord;

void dfs(int v, vector<vector<int>> &g, vector<bool> &used) {
    used[v] = true;
    for (int u : g[v]) {
        if (!used[u]) {
            dfs(u, g, used);
        }
    }
    ord.push_back(v);
}

void dfs(int v, vector<vector<int>> &g, vector<int> &color, int cur) {
    color[v] = cur;
    for (int u : g[v]) {
        if (!color[u]) {
            dfs(u, g, color, cur);
        }
    }
}

void solve() {
    int n, k;
    cin >> n >> k;
    vector<vector<int>> g(n);
    for (int i = 1; i < n; ++i) {
        int u, v;
        cin >> u >> v;
        --u; --v;
        g[u].push_back(v);
        g[v].push_back(u);
    }
    fill(par, par + maxn, -1);
    vector<int> c(n);
    for (int& i : c) cin >> i;
    for (int& i : c) --i;
    dfs(0, g, 0);
    vector<vector<int>> gp(k);
    vector<vector<int>> rgp(k);
    vector<vector<int>> groups(k);
    for (int i = 0; i < n; ++i) {
        groups[c[i]].push_back(i);
    }
    for (vector<int> g : groups) {
        vector<pair<int, int>> x;
        for (int i : g) {
            x.push_back({ tin[i], i });
        }
        sort(x.begin(), x.end());
        vector<pair<int, int>> y = x;
        for (int i = 0; i + 1 < x.size(); ++i) {
            y.push_back({ tin[lca(x[i].second, x[i + 1].second)], lca(x[i].second, x[i + 1].second) });
        }
        sort(y.begin(), y.end());
        vector<int> st;
        for (auto [x, v] : y) {
            while (!st.empty() && !isa(st.back(), v)) st.pop_back();
            if (!st.empty()) {
                int u = v;
                while (u != st.back()) {
                    gp[c[v]].push_back(c[u]);
                    rgp[c[u]].push_back(c[v]);
                    u = par[u];
                }
                gp[c[v]].push_back(c[u]);
                rgp[c[u]].push_back(c[v]);
            }
            st.push_back(v);
        }
    }
    vector<bool> used(k, false);
    for (int i = 0; i < k; ++i) {
        if (!used[i]) dfs(i, gp, used);
    }
    reverse(ord.begin(), ord.end());
    vector<int> color(k, 0);
    int cur = 1;
    for (int i : ord) {
        if (!color[i]) {
            dfs(i, rgp, color, cur++);
        }
    }
    vector<int> sz(k, 1);
    vector<int> sum(k, 0);
    for (int i = 0; i < k; ++i) {
        sum[color[i] - 1] += sz[i];
    }
    int res = inf;
    --cur;
    vector<bool> ok(cur, true);
    for (int i = 0; i < k; ++i) {
        for (int j : gp[i]) {
            if (color[j] != color[i]) {
                ok[color[i] - 1] = false;
            }
        }
    }
    for (int i = 0; i < cur; ++i) {
        if (ok[i]) {
            res = min(res, sum[i]);
        }
    }
    cout << res - 1;
}

int32_t main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
    int t;
    t = 1;
    while (t--) {
        solve();
    }
}
/*
 12 4
 7 9
 1 3
 4 6
 2 4
 10 12
 1 2
 2 10
 11 1
 2 8
 5 3
 6 7
 3
 1 1 2 4 3 3 2 2 3 4 4
 */

Compilation message (stderr)

capital_city.cpp: In function 'void solve()':
capital_city.cpp:106:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::pair<int, int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  106 |         for (int i = 0; i + 1 < x.size(); ++i) {
      |                         ~~~~~~^~~~~~~~~~

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