Submission #1088945

#	Time	Username	Problem	Language	Result	Execution time	Memory
1088945		_callmelucian	Capital City (JOI20_capital_city)	C++14	100 / 100	1748 ms	456120 KiB

capital_city

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

typedef long long ll;
typedef long double ld;
typedef pair<ll,ll> pl;
typedef pair<int,int> pii;
typedef tuple<int,int,int> tt;

#define all(a) a.begin(), a.end()
#define filter(a) a.erase(unique(all(a)), a.end())

const int mn = 2e5 + 5;
const int LOG = 18;

int go[mn][LOG], num[mn], depth[mn], clr[mn], sz[mn], timeDfs, n, k;
vector<int> nodeList[mn], adj[mn];

namespace SCC {
    const int NODE = 3e6 + 8e5 + 6;
    int num[NODE], low[NODE], gr[NODE], ans[NODE], timeDfs, counter = 0;
    vector<int> graph[NODE];
    bool del[NODE];
    stack<int> st;

    void dfs (int u, int p) {
        num[u] = low[u] = ++timeDfs;
        st.push(u);
        for (int v : graph[u]) {
            if (del[v]) continue;
            if (!num[v]) {
                dfs(v, u);
                low[u] = min(low[u], low[v]);
            }
            else low[u] = min(low[u], num[v]);
        }

        if (low[u] == num[u]) {
            int cur = 0, v = 0, flag = 1;
            counter++;
            do {
                v = st.top(); st.pop();
                del[v] = 1, cur += (v > LOG * n), gr[v] = counter;
                for (int oth : graph[v])
                    if (gr[oth] && gr[oth] != counter) flag = 0;
                if (!v) flag = 0;
            } while (v != u);
            ans[counter] = (flag ? cur : INT_MAX);
        }
    }

    int compress() {
        for (int i = 1; i <= LOG * n + k; i++)
            if (!num[i]) dfs(i, i);

        int res = INT_MAX;
        for (int i = 1; i <= counter; i++) res = min(res, ans[i]);
        return res;
    }
};

int node (int layer, int u) {
    return layer * n + u;
}

void addEdge (int a, int b) {
    SCC::graph[a].push_back(b);
}

void dfs (int u, int p, int d) {
    go[u][0] = p, depth[u] = d, num[u] = ++timeDfs, sz[u] = 1;
    if (p) addEdge(node(0, u), node(LOG, clr[p]));
    else addEdge(node(0, u), 0);

    for (int i = 1; i < LOG; i++) {
        go[u][i] = go[go[u][i - 1]][i - 1];
        addEdge(node(i, u), node(i - 1, u));
        if (go[u][i - 1]) addEdge(node(i, u), node(i - 1, go[u][i - 1]));
    }

    for (int v : adj[u]) {
        if (v == p) continue;
        dfs(v, u, d + 1);
        sz[u] += sz[v];
    }
}

int goUp (int a, int k) {
    for (int i = 0; i < LOG; i++)
        if (k & (1 << i)) a = go[a][i];
    return a;
}

int lca (int a, int b) {
    if (depth[a] < depth[b]) swap(a, b);
    a = goUp(a, depth[a] - depth[b]);
    if (a == b) return a;
    for (int i = LOG - 1; i >= 0; i--)
        if (go[a][i] != go[b][i]) a = go[a][i], b = go[b][i];
    return go[a][0];
}

void connectPath (int source, int a, int p) {
    //cout << "connecting " << source << " " << a << " " << p << "\n";
    addEdge(source, node(LOG, clr[a]));
    for (int i = LOG - 1; i >= 0; i--) {
        if (depth[go[a][i]] < depth[p]) continue;
        addEdge(source, node(i, a));
        a = go[a][i];
    }
}

bool cmp (int a, int b) {
    return num[a] < num[b];
}

bool isParent (int p, int u) {
    return num[p] <= num[u] && num[u] < num[p] + sz[p];
}

int main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);

    cin >> n >> k;
    for (int i = 1; i < n; i++) {
        int a, b; cin >> a >> b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    for (int i = 1; i <= n; i++) {
        cin >> clr[i];
        nodeList[clr[i]].push_back(i);
    }
    dfs(1, 0, 1);

    for (int color = 1; color <= k; color++) {
        if (nodeList[color].empty()) continue;
        sort(all(nodeList[color]), cmp);
        int curSz = nodeList[color].size();
        for (int i = 1; i < curSz; i++)
            nodeList[color].push_back(lca(nodeList[color][i - 1], nodeList[color][i]));
        sort(all(nodeList[color]), cmp);
        filter(nodeList[color]);

        stack<int> st({nodeList[color][0]});
        for (int i = 1; i < nodeList[color].size(); i++) {
            while (st.size() && !isParent(st.top(), nodeList[color][i])) st.pop();
            if (st.size())
                connectPath(node(LOG, color), nodeList[color][i], st.top());
            st.push(nodeList[color][i]);
        }
    }

    cout << SCC::compress() - 1;

    return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'int main()':
capital_city.cpp:148:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  148 |         for (int i = 1; i < nodeList[color].size(); i++) {
      |                         ~~^~~~~~~~~~~~~~~~~~~~~~~~

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