Submission #1322297

#	Time	Username	Problem	Language	Result	Execution time	Memory
1322297		phungmanager0	Split the Attractions (IOI19_split)	C++20	100 / 100	92 ms	27140 KiB

split

#include <bits/stdc++.h>
#include "split.h"
#define FOR(i, a, b) for (int i = (a), _b = (b); i <= _b; i++)
#define FORD(i, b, a) for (int i = (b), _a = (a); i >= _a; i--)
#define REP(i, n) for (int i = 0, _n = (n); i < _n; i++)
#define FORE(i, v) for (__typeof((v).begin()) i = (v).begin(); i != (v).end(); i++)
#define ALL(v) (v).begin(), (v).end()
using namespace std;

struct DisjointSet {
    vector<int> par, Sz;
    void init(int n) { par.resize(n + 5); Sz.resize(n + 5); 
REP(i, n) par[i] = i, Sz[i] = 1;}
    int get(int u) { return (par[u] == u ? u : par[u] = get(par[u])); }
    bool merge(int u, int v) {
        u = get(u); v = get(v); 
        if(u == v) return false;
        if(Sz[u] < Sz[v]) swap(u, v);
        par[v] = u; Sz[u] += Sz[v];
        return true;
    }
};
struct updateDisjointSet {
    vector<int> par, Sz, total_val; void init(int n) { par.resize(n + 5); Sz.resize(n + 5); total_val.resize(n + 5); 
REP(i, n) par[i] = i, Sz[i] = 1; }
    void update(int idx, int W) { total_val[idx] = W; }
    int get(int u) { return (u == par[u] ? u : par[u] = get(par[u])); }
    void merge(int u, int v) {
        u = get(u); v = get(v); if(u == v) return; if(Sz[u] < Sz[v]) swap(u, v);
        par[v] = u; Sz[u] += Sz[v]; total_val[u] += total_val[v];
    }
};
bool subtaks1(vector<int> p, vector<int> q, int n) {
    vector<int> cnt(n); REP(i, (int)q.size()) { cnt[p[i]]++; cnt[q[i]]++; }
    REP(i, n) if(cnt[i] >= 3) return false;
    return true;
}
bool cmp(const pair<int, int>& a, const pair<int, int>& b) { return a.first < b.first; }
vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    vector<pair<int, int>> cmp_list = {{a, 1}, {b, 2}, {c, 3}}; sort(ALL(cmp_list), cmp);
    vector<int> res(n); if(subtaks1(p, q, n)) {
        DisjointSet dsu; dsu.init(n); vector<vector<int>> g; g.resize(n + 5);
        REP(i, (int)p.size()) if(dsu.merge(p[i], q[i])) {
            g[p[i]].push_back(q[i]); g[q[i]].push_back(p[i]);
        } vector<int> h(n);
        function<void(int, int)> DFS = [&](int u, int p) {
            FORE(it, g[u]) if((*it) != p) {
                h[(*it)] = h[u] + 1; DFS(*it, u);
            }
        };
        DFS(0, -1);
        int newRoot = 0; FOR(i, 1, n - 1) if(h[i] > h[newRoot]) newRoot = i;
        vector<int> child(n);
        function<void(int, int)> DFS_Child = [&](int u, int p) {
            int next_v = -1;
            FORE(it, g[u]) if((*it) != p) { next_v = (*it); break; }
            if(next_v != -1) {
                child[u] = next_v; DFS_Child(child[u], u);
            }
        };
        DFS_Child(newRoot, -1); int cur_cnt = 0;
        while(cur_cnt < a) {
            cur_cnt++; res[newRoot] = 1; newRoot = child[newRoot];
        }
        cur_cnt = 0;
        while(cur_cnt < b) {
            cur_cnt++; res[newRoot] = 2; newRoot = child[newRoot];
        }
        REP(i, n) if(res[i] == 0) res[i] = 3;
    } else if(a == 1) {
        DisjointSet dsu; dsu.init(n);
        vector<vector<int>> g; g.resize(n + 5); sort(ALL(cmp_list), cmp);
        REP(i, (int)p.size()) {
            g[p[i]].push_back(q[i]); g[q[i]].push_back(p[i]);
            dsu.merge(p[i], q[i]);
        }
        bool _check = false; int idx = -1;
        REP(i, n) if(dsu.Sz[dsu.get(i)] >= cmp_list[1].first + 1) { _check = true; idx = i; break; }
        if(_check == false) return res;
        int cnt = 0; vector<bool> check(n); check[idx] = true;
        function<void(int)> DFS = [&](int u) {
            if(cnt == cmp_list[1].first - 1) return;
            FORE(it, g[u]) if(!check[*it]) {
                if(cnt == cmp_list[1].first - 1) return;
                res[*it] = cmp_list[1].second; cnt++; check[*it] = true; 
                DFS(*it); if(cnt == cmp_list[1].first - 1) return;

            }
        }; DFS(idx); res[idx] = cmp_list[1].second; 
        REP(i, n) if(res[i] != cmp_list[1].second && dsu.get(i) == dsu.get(idx)) 
        { res[i] = cmp_list[0].second; break; }
        REP(i, n) if(res[i] == 0) res[i] = cmp_list[2].second;
    } else if((int)p.size() == n - 1) {
        vector<vector<int>> g; g.resize(n + 5); vector<int> Sz(n), par(n);
        REP(i, (int)p.size()) { g[p[i]].push_back(q[i]); g[q[i]].push_back(p[i]); }
        function<void(int, int)> DFS = [&](int u, int p) {
            Sz[u] = 1; FORE(it, g[u]) if((*it) != p) {
                par[(*it)] = u;
                DFS(*it, u); Sz[u] += Sz[*it];
            }
        };
        function<int(int, int)> FindCentroid = [&](int u, int p) {
            FORE(it, g[u]) if((*it) != p && Sz[*it] > n / 2) {
                return FindCentroid(*it, u);
            }
            return u;
        }; DFS(1, -1); int centroid = FindCentroid(1, -1); vector<pair<int, int>> comp_size;
        FORE(it, g[centroid]) if((*it) != par[centroid]) comp_size.push_back({Sz[(*it)], (*it)});
        int add = n - 1, idx = -1; FORE(it, comp_size) add -= (*it).first; comp_size.push_back({add, par[centroid]});
        bool check = false; FORE(it, comp_size) if((*it).first >= cmp_list[0].first) {
            check = true; idx = (*it).second;
        }
        int cnt = 0;
        if(!check) return res;
        function<void(int, int)> DFS_Fill = [&](int u, int p) {
            res[u] = cmp_list[0].second; cnt++; if(cnt == cmp_list[0].first) return;
            FORE(it, g[u]) if((*it) != p) {
                if(cnt == cmp_list[0].first) return;
                DFS_Fill(*it, u);
                if(cnt == cmp_list[0].first) return;
            }
        };
        DFS_Fill(idx, centroid); 
        queue<int> q; q.push(centroid); cnt = 0; while(q.size()) {
            int u = q.front(); q.pop(); cnt++; res[u] = cmp_list[1].second;
            if(cnt == cmp_list[1].first) break;
            FORE(it, g[u]) if(res[*it] == 0) {
                q.push(*it);
            }
        }
        REP(i, n) if(res[i] == 0) res[i] = cmp_list[2].second;
    } else {
        DisjointSet dsu, compress; dsu.init(n); compress.init(n);
        vector<int> h(n), Sz(n); vector<pair<int, int>> edges, edgesspanningtree, otheredges; 
        REP(i, (int)p.size()) {
            edges.push_back({p[i], q[i]});
        }
        vector<vector<int>> spanningtree, g; spanningtree.resize(n + 5); g.resize(n + 5); int u = 0, v = 0;
        FORE(it, edges) {
            u = (*it).first; v = (*it).second; if(dsu.merge(u, v)) {
                spanningtree[u].push_back(v); spanningtree[v].push_back(u);
                edgesspanningtree.push_back({u, v});
            } else otheredges.push_back({u, v});
            g[u].push_back(v); g[v].push_back(u);
        }
        function<void(int, int)> DFS = [&](int u, int p) {
            Sz[u] = 1;
            FORE(it, spanningtree[u]) if((*it) != p) {
                h[(*it)] = h[u] + 1; DFS((*it), u); Sz[u] += Sz[*it];
            }
        };
        function<int(int, int)> FindCentroid = [&](int u, int p) {
            FORE(it, spanningtree[u]) if((*it) != p && Sz[(*it)] > n / 2) {
                return FindCentroid(*it, u);
            }
            return u;
        };
        DFS(0, -1); int centroid = FindCentroid(0, -1);
        FORE(it, edgesspanningtree) {
            if((*it).first != centroid && (*it).second != centroid) {
                compress.merge((*it).first, (*it).second);
            }
        }
        vector<pair<int, int>> points; REP(i, n) if(i == compress.get(i)) {
            points.push_back({i, 0});
        }
        for(auto &[x, y] : points) y = compress.Sz[compress.get(x)]; int cnt = 0;
        function<void(int, int)> DFS_Fill = [&](int u, int p) {
            if(cnt == cmp_list[0].first) return;
            res[u] = cmp_list[0].second; cnt++; if(cnt == cmp_list[0].first) return;
            FORE(it, spanningtree[u]) if((*it) != p) {
                if(cnt == cmp_list[0].first) return;
                DFS_Fill(*it, u);
                if(cnt== cmp_list[0].first) return;
            }
        };
        bool check = false; int idx_comp = -1;
        FORE(it, points) {
            if((*it).second >= cmp_list[0].first && (*it).first != centroid) {
                check = true; idx_comp = (*it).first; break; }
        }
        if(check) {
            int next_idx = -1; 
            FORE(it, spanningtree[centroid]) if(compress.get(*it) == idx_comp) { next_idx = (*it); break; }
            else if((*it) == idx_comp) { next_idx = idx_comp; break; }
            if(next_idx == -1) return res;
            DFS_Fill(next_idx, centroid); queue<int> q; q.push(centroid); cnt = 0; while(q.size()) {
                int u = q.front(); q.pop(); cnt++; res[u] = cmp_list[1].second;
                if(cnt == cmp_list[1].first) break;
                FORE(it, spanningtree[u]) if(res[*it] == 0) {
                    q.push(*it);
                }
            }
            REP(i, n) if(res[i] == 0) res[i] = cmp_list[2].second;
        } else {
            updateDisjointSet compute; compute.init(n);
            REP(i, (int)points.size()) compute.update(points[i].first, points[i].second);
            vector<vector<int>> newGraph(n + 1); vector<int> Weight(n + 1);
            FORE(it, points) Weight[(*it).first] = (*it).second;
            int first_comp = -1, second_comp = -1;
            FORE(it, otheredges) {
                if((*it).first == centroid || (*it).second == centroid) continue;
                first_comp = compress.get((*it).first); second_comp = compress.get((*it).second);
                if(first_comp != second_comp) {
                    newGraph[first_comp].push_back(second_comp); newGraph[second_comp].push_back(first_comp);
                    compute.merge(first_comp, second_comp);
                }
            }
            bool check_ans = false; int start_bfs = 0; vector<bool> check(n + 1), allow(n + 1);
            FORE(it, points) if(compute.total_val[compute.get((*it).first)] >= cmp_list[0].first) { check_ans = true; start_bfs = (*it).first; break; }
            if(check_ans) {
                queue<int> q; int sumWeight = 0;
                q.push(start_bfs); while(q.size()) {
                    int u = q.front(); q.pop(); if(check[u]) continue; check[u] = true;
                    sumWeight += Weight[u]; if(sumWeight >= cmp_list[0].first) break;
                    else {
                        FORE(it, newGraph[u]) if(!check[(*it)]) q.push(*it);
                    }
                }
                REP(i, n) if(check[compress.get(i)]) allow[i] = true;
                while(q.size()) q.pop(); q.push(start_bfs); int cnt = 0; while(q.size()) {
                    int u = q.front(); q.pop();
                    if(res[u] != 0) continue; res[u] = cmp_list[0].second; cnt++; if(cnt == cmp_list[0].first) break;
                    else {
                        FORE(it, g[u]) if(res[*it] == 0 && allow[*it]) q.push(*it);
                    }
                } while(q.size()) q.pop(); q.push(centroid); cnt = 0; while(q.size()) {
                    int u = q.front(); q.pop(); if(res[u] != 0) continue; cnt++; res[u] = cmp_list[1].second;
                    if(cnt == cmp_list[1].first) break;
                    FORE(it, spanningtree[u]) if(res[*it] == 0) {
                        q.push(*it);
                    }
                }
                REP(i, n) if(res[i] == 0) res[i] = cmp_list[2].second;
            }
        }
    }
    return res;
}

• Subtask 1: 7 / 7
• Subtask 2: 11 / 11
• Subtask 3: 22 / 22
• Subtask 4: 24 / 24
• Subtask 5: 36 / 36