Submission #1221697

#TimeUsernameProblemLanguageResultExecution timeMemory
1221697turskaSplit the Attractions (IOI19_split)C++20
40 / 100
87 ms17724 KiB
#include "split.h"
#include <bits/stdc++.h>

using namespace std;
const int mxN = 1e5;
vector<int> t_adj[mxN];
int sz[mxN];

#define all(v) (v).begin(), (v).end()

mt19937_64 gen(chrono::high_resolution_clock::now().time_since_epoch().count());

vector<array<int, 2>> dfs_edges; // parent child
void dfs(int v, int p) {
    sz[v] = 1;
    for (auto u: t_adj[v]) if (u!=p) {
        dfs_edges.push_back({v, u});
        dfs(u, v);
        sz[v] += sz[u];
    }
}

void mark(int v, int p, int cnt, int clr, vector<int> &res) {
    queue<array<int, 2>> q;
    q.push({v, p});
    while (!q.empty() && cnt) {
        auto [u, p] = q.front(); q.pop();
        res[u] = clr;
        cnt--;
        for (auto w: t_adj[u]) if (w!=p) {
            q.push({w, u});
        }
    }
}
vector<int> try_find(int n, array<int, 3> sets, array<int, 3> ord) {
    vector<int> res(n);
    int x = sets[ord[0]];
    int y = sets[ord[1]];
    for (auto [p, c]: dfs_edges) {
        int sz_c = sz[c];
        int sz_p = n-sz[c];
        if (sz_c>sz_p) {
            swap(c, p);
            swap(sz_c, sz_p);
        }
        if (sz_c>= x && sz_p >= y)  {
            mark(c, p, x, ord[0]+1, res);
            mark(p, c, y, ord[1]+1, res);
            for (int i=0; i<n; i++) if (!res[i]) res[i] = ord[2]+1;
            break;
        }
    }
    return res;
}
struct DSU {
    vector<int> p, sz;
    DSU(int n): p(n, 0), sz(n, 0) {
        iota(all(p), 0);
    }
    int find(int v) {
        if (p[v]==v) return v; 
        return p[v] = find(p[v]);
    }
    bool unite(int a, int b) {
        a = find(a);
        b = find(b);
        if (a==b) return false;
        if (sz[a]<sz[b]) swap(a, b);
        p[b] = a;
        sz[a] += sz[b];
        return true;
    }
};
void gen_spanning_tree(int n, vector<array<int, 2>> &e) {
    DSU dsu(n);
    shuffle(all(e), gen);
    for (int i=0; i<n; i++) t_adj[i].clear();

    for (auto [u, v]: e) {
        if (dsu.unite(u, v)) {
            t_adj[u].push_back(v);
            t_adj[v].push_back(u);
        }
    }
    dfs_edges.clear();
    fill(sz, sz+n, 0);
    dfs(0, -1);
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    vector<array<int, 2>> e;
    for (int i=0; i<p.size(); i++) {
        int u = p[i];
        int v = q[i];
        e.push_back({u, v});
    }
    array<int, 3> sets = {a, b, c};
    array<int, 3> ord = {0, 1, 2};
    sort(all(ord), [&](int l, int r) {
        return sets[l] < sets[r];
    });
    gen_spanning_tree(n, e);
    return try_find(n, sets, ord);
}
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...