Submission #900136

#	Time	Username	Problem	Language	Result	Execution time	Memory
900136		DAleksa	Capital City (JOI20_capital_city)	C++17	100 / 100	183 ms	71232 KiB

capital_city

#include <bits/stdc++.h>

using namespace std;

const int N = 2e5 + 10, LOG = 19;
int n, k;
vector<int> tree[N];
vector<int> g[N], rg[N];
int a[N];
vector<int> pos[N];
int up[N][LOG];
int tin[N], tout[N], timer;
int lca[N];
bool mark[N];
bool outdeg[N];
int sz[N];

void dfs(int u, int par) {
    tin[u] = timer++;
    for(int v : tree[u]) {
        if(v == par) continue;
        up[v][0] = u;
        dfs(v, u);
    }
    tout[u] = timer++;
}

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

int LCA(int u, int v) {
    if(is_ancestor(u, v)) return u;
    if(is_ancestor(v, u)) return v;
    for(int j = LOG - 1; j >= 0; j--) {
        if(up[u][j] == 0) continue;
        if(is_ancestor(up[u][j], v)) continue;
        u = up[u][j];
    }
    return up[u][0];
}

vector<int> vec;
void dfsrg(int u) {
    if(mark[u]) return;
    mark[u] = true;
    for(int v : rg[u]) {
        if(mark[v]) continue;
        dfsrg(v);
    }
    vec.push_back(u);
}

int scc[N];
int id = 0;
void dfsg(int u) {
    if(mark[u]) return;
    mark[u] = true;
    scc[u] = id;
    for(int v : g[u]) {
        if(mark[v]) continue;
        dfsg(v);
    }
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cin >> n >> k;
    for(int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }
    for(int i = 1; i <= n; i++) {
        cin >> a[i];
        pos[a[i]].push_back(i);
    }
    dfs(1, 0);
    for(int j = 1; j < LOG; j++) for(int i = 1; i <= n; i++) up[i][j] = up[up[i][j - 1]][j - 1];
    for(int i = 1; i <= k; i++) {
        lca[i] = pos[i][0];
        for(int j = 1; j < pos[i].size(); j++) lca[i] = LCA(lca[i], pos[i][j]);
    }
    for(int i = 2; i <= n; i++) {
        if(i == lca[a[i]]) continue;
        if(a[up[i][0]] != a[i]) g[a[i]].push_back(a[up[i][0]]);
    }
    for(int i = 1; i <= k; i++) {
        sort(g[i].begin(), g[i].end());
        g[i].erase(unique(g[i].begin(), g[i].end()), g[i].end());
        for(int j : g[i]) rg[j].push_back(i);
    }
    for(int i = 1; i <= k; i++) {
        if(mark[i]) continue;
        dfsrg(i);
    }
    reverse(vec.begin(), vec.end());
    for(int i = 1; i <= k; i++) mark[i] = false;
    for(int i : vec) {
        if(mark[i]) continue;
        id++;
        dfsg(i);
    }
    for(int i = 1; i <= k; i++) {
        for(int j : g[i]) if(scc[i] != scc[j]) outdeg[scc[i]] = true;
        sz[scc[i]]++;
    }
    int ans = k - 1;
    for(int i = 1; i <= k; i++) if(!outdeg[scc[i]]) ans = min(ans, sz[scc[i]] - 1);
    cout << ans;
    return 0;
}

Compilation message (stderr)

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

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