Submission #143590

#	Time	Username	Problem	Language	Result	Execution time	Memory
143590		VladaMG98	Split the Attractions (IOI19_split)	C++17	18 / 100	124 ms	20068 KiB

split

#include <bits/stdc++.h>
#include "split.h"
using namespace std;

const int MAXN = 100010;
int subtree_size[MAXN];
vector<int> adj[MAXN];
int parent[MAXN];

void dfs(int src = 0, int prev = -1) {
    //printf("dfs %d %d\n", src, prev);
    subtree_size[src] = 1;
    parent[src] = prev;
    for(auto &xt : adj[src]) {
        if(xt - prev) {
            dfs(xt, src);
            subtree_size[src] += subtree_size[xt];
        }
    }
}

vector<int> get_solution(vector<int> v1, int n1, vector<int> v2, int n2, int n) {
    vector<int> ans(n, 6 - n1 - n2);
    for(auto &x : v1) {
        ans[x] = n1;
    }
    for(auto &x : v2) {
        ans[x] = n2;
    }
    return ans;
}

bool marked[MAXN];
vector<int> extract(int src, int wrong, int sz) {
    marked[wrong] = true;
    vector<int> ans;
    stack<int> nodes;
    nodes.push(src);
    marked[src] = true;
    while((int)ans.size() < sz) {
        int node = nodes.top();
        nodes.pop();
        ans.push_back(node);
        for(auto &x : adj[node]) {
            if(!marked[x]) {
                marked[x] = true;
                nodes.push(x);
            }
        }
    }
    for(auto &x : ans) {
        marked[x] = false;
    }
    marked[wrong] = false;
    return ans;
}

vector<int> blank(int n) {
    vector<int> ans(n, 0);
    return ans;
}

int dsu[MAXN];

void init_set(int x) {
    dsu[x] = x;
}

int get(int x) {
    if(dsu[x] == x) return x;
    return dsu[x] = get(dsu[x]);
}

bool unite(int x, int y) {
    int xr = get(x), yr = get(y);
    if(xr == yr) return false;
    if(rand() % 2) swap(xr, yr);
    dsu[xr] = yr;
    return true;
}

void create_spanning_tree(int n, vector<int> &st_p, vector<int> &st_q, vector<int> p, vector<int> q) {
    for(int i = 0; i < n; i++) {
        init_set(i);
    }
    int m = (int)p.size();
    for(int i = 0; i < m; i++) {
        if(unite(p[i], q[i])) {
            st_p.push_back(p[i]);
            st_q.push_back(q[i]);
        }
    }
    //assert((int)st_p.size() == n - 1);
}

int belongs[MAXN];
set<int> comp_adj[MAXN];
int cnt_nodes[MAXN];

int comp_size;
vector<int> comp_comp;
bool comp_mark[MAXN];
void dfs2(int src) {
    comp_size += cnt_nodes[src];
    comp_mark[src] = true;
    comp_comp.push_back(src);
    for(auto &xt : comp_adj[src]) {
        if(!comp_mark[xt]) {
            dfs2(xt);
        }
    }
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    pair<int, int> arr[] = {{a, 1}, {b, 2}, {c, 3}};
    sort(arr, arr + 3);
    int A = arr[0].first, B = arr[1].first, C = arr[2].first;
    vector<int> st_p, st_q;
    create_spanning_tree(n, st_p, st_q, p, q);
    int m = (int)p.size();
    for(int i = 0; i < n - 1; i++) {
        adj[st_p[i]].push_back(st_q[i]);
        adj[st_q[i]].push_back(st_p[i]);
    }
    dfs();
    //for(int i = 0; i < n; i++) {
    //   printf("subtree_size[%d] = %d\n", i, subtree_size[i]);
    //}
    for(int i = 0; i < m; i++) {
        int u = p[i], v = q[i];
        if(subtree_size[u] > subtree_size[v]) swap(u, v);
        int under = subtree_size[u];
        int rest = n - under;
        if(under >= A && rest >= A) {
            //printf("found for edge %d - %d\n", u, v);
            //we've got a solution
            vector<int> set_A, set_B;
            if(under < rest) {
                set_A = extract(u, v, A);
                set_B = extract(v, u, B);
            } else {
                set_A = extract(v, u, A);
                set_B = extract(u, v, B);
            }
            return get_solution(set_A, arr[0].second, set_B, arr[1].second, n);
        }
    }
    //there is no edge connecting two subtrees of sizes >= a
    //meaning that it's a star with all subtrees attached of size < a
    //in tree, there is no such solution
    for(int node = 0; node < n; node++) {
        vector<pair<int, int>> sizes;
        int sm = 0;
        for(auto &xt : adj[node]) {
            if(xt != parent[node]) {
                sizes.push_back({subtree_size[xt], xt});
                sm += subtree_size[xt];
            }
        }
        if(node) {
            sizes.push_back({n - 1 - sm, parent[node]});
        }
        int mx = 0;
        for(auto &pr : sizes) {
            mx = max(mx, pr.second);
        }
        if(mx < A) {
            printf("I am deciding on node %d\n", node);
            for(int i = 0; i < (int)sizes.size(); i++) {
                vector<int> comp_here = extract(sizes[i].second, node, sizes[i].first);
                for(auto &x : comp_here) {
                    belongs[x] = i + 1;
                    marked[x] = false;
                }
                marked[node] = false;
                cnt_nodes[i + 1] = sizes[i].first;
            }
            int new_nodes = (int)sizes.size();
            for(int i = 0; i < m; i++) {
                if(p[i] != node && q[i] != node) {
                    comp_adj[belongs[p[i]]].insert(belongs[q[i]]);
                    comp_adj[belongs[q[i]]].insert(belongs[p[i]]);
                }
            }
            for(int i = 1; i <= new_nodes; i++) {
                if(marked[i]) continue;
                comp_size = 0;
                comp_comp.clear();
                dfs2(i);
                for(auto &x : comp_comp) {
                    comp_mark[x] = false;
                }
                if(comp_size >= A) {
                    //we have a solution
                    int start_node = sizes[i - 1].second;
                    vector<int> ans_A = extract(start_node, node, A);
                    for(auto &x : ans_A) {
                        marked[x] = true;
                    }
                    vector<int> ans_B = extract(node, MAXN - 1, B);
                    return get_solution(ans_A, arr[0].second, ans_B, arr[1].second, n);
                } else {
                    return blank(n);
                }
            }
        }
    }
    assert(false);
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:117:45: warning: unused variable 'C' [-Wunused-variable]
     int A = arr[0].first, B = arr[1].first, C = arr[2].first;
                                             ^

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