Submission #1321199

#	Time	Username	Problem	Language	Result	Execution time	Memory
1321199		phungmanager0	Split the Attractions (IOI19_split)	C++20	40 / 100	67 ms	18180 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;
    }
};
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); // bfs form 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;
    }
    return res;
}

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