Submission #1036728

#	Time	Username	Problem	Language	Result	Execution time	Memory
1036728		vladilius	Capital City (JOI20_capital_city)	C++17	100 / 100	552 ms	157544 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;
using ll = long long;
using pii = pair<int, int>;
#define pb push_back
#define ff first
#define ss second
#define ins insert

struct STG{
    vector<vector<int>> g;
    vector<int> t;
    int n, k;
    STG(int& ns, int& ks, vector<int>& ts){
        n = ns;
        k = ks;
        t = ts;
        g.resize(4 * n + k);
        build(1, 1, n);
    }
    void build(int v, int tl, int tr){
        if (tl == tr){
            g[v + k].pb(t[tl]);
            return;
        }
        int tm = (tl + tr) / 2, vv = 2 * v;
        build(vv, tl, tm);
        build(vv + 1, tm + 1, tr);
        v += k; vv += k;
        g[v].pb(vv); g[v].pb(vv + 1);
    }
    void add(int v, int tl, int tr, int& l, int& r, int& u){
        if (l > tr || r < tl) return;
        if (l <= tl && tr <= r){
            g[u].pb(v + k);
            return;
        }
        int tm = (tl + tr) / 2, vv = 2 * v;
        add(vv, tl, tm, l, r, u);
        add(vv + 1, tm + 1, tr, l, r, u);
    }
    void add(int l, int r, int u){
        add(1, 1, n, l, r, u);
    }
    vector<vector<int>> gt;
    vector<int> order;
    vector<bool> used;
    void dfs1(int v){
        used[v] = 1;
        for (int i: g[v]){
            if (!used[i]){
                dfs1(i);
            }
        }
        order.pb(v);
    }
    vector<int> cp, cc;
    int cur;
    void dfs2(int v){
        cp[v] = cur;
        cc[cur] += (v <= k);
        for (int i: gt[v]){
            if (!cp[i]){
                dfs2(i);
            }
        }
    }
    vector<vector<int>> gg;
    void dfs3(int v){
        used[v] = 1;
        for (int i: gg[v]){
            if (!used[i]){
                dfs3(i);
            }
        }
    }
    void scc(){
        used.resize(g.size());
        for (int i = 1; i <= k + 1; i++){
            if (!used[i]) dfs1(i);
        }
        reverse(order.begin(), order.end());
        
        gt.resize(g.size());
        for (int i = 1; i < g.size(); i++){
            for (int j: g[i]){
                gt[j].pb(i);
            }
        }
        
        cp.resize(g.size());
        cc.resize(g.size());
        cur = 0;
        for (int i: order){
            if (!cp[i]){
                cur++;
                dfs2(i);
            }
        }
        
        gg.resize(cur + 1);
        for (int i = 1; i < g.size(); i++){
            for (int j: g[i]){
                if (cp[i] != cp[j]){
                    gg[cp[i]].pb(cp[j]);
                }
            }
        }
        int out = k;
        for (int i = 1; i <= cur; i++){
            if (gg[i].empty()){
                out = min(out, cc[i]);
            }
        }
        cout<<--out<<"\n";
    }
};

int main(){
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(0);
    
    int n, k; cin>>n>>k;
    vector<int> g[n + 1];
    for (int i = 1; i < n; i++){
        int x, y; cin>>x>>y;
        g[x].pb(y);
        g[y].pb(x);
    }
    vector<int> ct[k + 1], c(n + 1);
    for (int i = 1; i <= n; i++){
        cin>>c[i];
        ct[c[i]].pb(i);
    }
    
    vector<int> sz(n + 1), p(n + 1), h(n + 1), d(n + 1);
    function<void(int, int)> fill = [&](int v, int pr){
        sz[v] = 1;
        p[v] = pr;
        for (int i: g[v]){
            if (i == pr) continue;
            d[i] = d[v] + 1;
            fill(i, v);
            sz[v] += sz[i];
            if (sz[i] > sz[h[v]]){
                h[v] = i;
            }
        }
    };
    fill(1, 0);
    
    vector<int> head(n + 1), pos(n + 1);
    int timer = 0;
    function<void(int, int)> fill_hld = [&](int v, int H){
        head[v] = H;
        pos[v] = ++timer;
        if (!h[v]) return;
        fill_hld(h[v], H);
        for (int i: g[v]){
            if (pos[i]) continue;
            fill_hld(i, i);
        }
    };
    fill_hld(1, 1);
    
    vector<int> t(n + 1);
    for (int i = 1; i <= n; i++){
        t[pos[i]] = c[i];
    }
    
    STG T(n, k, t);
    
    auto add = [&](int x, int y, int i){
        while (head[x] != head[y]){
            if (d[head[x]] > d[head[y]]) swap(x, y);
            T.add(pos[head[y]], pos[y], i);
            y = p[head[y]];
        }
        if (d[x] > d[y]) swap(x, y);
        T.add(pos[x], pos[y], i);
    };
    
    for (int i = 1; i <= k; i++){
        for (int j = 0; j + 1 < ct[i].size(); j++){
            add(ct[i][j], ct[i][j + 1], i);
        }
    }
    
    T.scc();
}

Compilation message (stderr)

capital_city.cpp: In member function 'void STG::scc()':
capital_city.cpp:85:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   85 |         for (int i = 1; i < g.size(); i++){
      |                         ~~^~~~~~~~~~
capital_city.cpp:102:27: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<std::vector<int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  102 |         for (int i = 1; i < g.size(); i++){
      |                         ~~^~~~~~~~~~
capital_city.cpp: In function 'int main()':
capital_city.cpp:185:31: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  185 |         for (int j = 0; j + 1 < ct[i].size(); j++){
      |                         ~~~~~~^~~~~~~~~~~~~~

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