Submission #1237650

#	Time	Username	Problem	Language	Result	Execution time	Memory
1237650		jer033	Split the Attractions (IOI19_split)	C++20	40 / 100	99 ms	29684 KiB

split

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

vector<vector<int>> new_edges;

int centroid_find(int n, vector<vector<int>> edges)
{
    //create a spanning tree
    //vector<vector<int>> new_edges(n);
    vector<int> parent(n, -2);
    vector<int> nodes;
    new_edges = vector<vector<int>> (n);

    nodes.push_back(0);
    parent[0] = -1;
    
    int curr = 0;
    int total = 1;

    while (curr < total)
    {
        int x = nodes[curr]; curr++;
        for (int y: edges[x])
        {
            if (parent[y] == -2)
            {
                parent[y] = x;
                nodes.push_back(y);
                new_edges[x].push_back(y);
                new_edges[y].push_back(x);
                total++;
            }
        }
    }

    vector<int> subtreesize(n, 1);
    for (int i=n-1; i>0; i--)
    {
        int x = nodes[i];
        subtreesize[parent[x]] += subtreesize[x];
    }

    //now identify the centroid
    int bes = n+1;
    int cap = (n+1)/2;
    int centroid = -1;
    for (int i=0; i<n; i++)
        if ((subtreesize[i]<bes) and (subtreesize[i]>=cap))
        {
            bes = subtreesize[i];
            centroid = i;
        }
    return centroid;
}

vector<int> find_split_given_two_colors(int n, int a, int b, int c, vector<vector<int>> edges, vector<int> colors)
{
    int col_4 = 0;
    int col_5 = 0;
    for (int i=0; i<n; i++)
    {
        if (colors[i]==4)
            col_4++;
        else
            col_5++;
    }
    if (col_4 >= b)
    {
        for (int i=0; i<n; i++)
            colors[i] = 9-colors[i];
        swap(col_4, col_5);
    }

    int root, count;
    root = 0;
    count = 1;
    while (colors[root]!=4)
        root++;
    colors[root] = 1;
    queue<int> nodes4;
    nodes4.push(root);
    while (count < a)
    {
        int x = nodes4.front(); nodes4.pop();
        for (int y: edges[x])
            if ((colors[y] == 4) and (count < a))
            {
                colors[y] -= 3;
                count++;
                nodes4.push(y);
            }
    }

    root = 0;
    count = 1;
    while (colors[root]!=5)
        root++;
    colors[root] = 2;
    queue<int> nodes5;
    nodes5.push(root);
    while (count < b)
    {
        int x = nodes5.front(); nodes5.pop();
        for (int y: edges[x])
            if ((colors[y] == 5) and (count < b))
            {
                colors[y] -= 3;
                count++;
                nodes5.push(y);
            }
    }

    for (int i=0; i<n; i++)
        if (colors[i]>3)
            colors[i] = 3;

    return colors;
}

vector<int> find_split_wlog(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	vector<vector<int>> edges(n);
    int e = p.size();
    for (int i=0; i<e; i++)
    {
        int pp = p[i];
        int qq = q[i];
        edges[pp].push_back(qq);
        edges[qq].push_back(pp);
    }

    int root = centroid_find(n, edges);
    //cout << root << '\n';

    //step 3: perform a dfs and identify whether its possible or not

    vector<int> parent(n, -2);
    stack<pair<int, int>> nodes;
    vector<int> push_order;
    vector<vector<int>> children(n);

    parent[root] = -1;
    nodes.push({root, 0});
    push_order.push_back(root);

    while (!nodes.empty())
    {
        pair<int, int> info = nodes.top();
        int x = info.first; int y = info.second;
        nodes.pop();
        if ((y+1)<edges[x].size())
            nodes.push({x, y+1});
        int z = edges[x][y];
        if (parent[z]==-2)
        {
            parent[z] = x;
            nodes.push({z, 0});
            push_order.push_back(z);
            children[x].push_back(z);
        }
    }

    vector<int> subtreesize(n, 1);
    for (int i=n-1; i>0; i--)
    {
        int x = push_order[i];
        subtreesize[parent[x]] += subtreesize[x];
        //cout << x << subtreesize[x] << '\n';
    }
    
    //determine impossibility
    int node_min = n+1;
    int root_a = -1;
    for (int i: children[root])
        if ((i!=root) and (subtreesize[i] < node_min) and (subtreesize[i]>=a))
        {
            node_min = subtreesize[i];
            root_a = i;
        }
    

    if (root_a == -1)
    {
        vector<int> ans(n, 0);
        return ans;
    }

    //then a solution is possible
    
    //do stuff with root_a

    //cout << "reached here - 1\n";
    parent = vector<int>(n, -2);
    vector<int> color(n, -1);
    vector<int> color_count(n, 0);
    queue<int> nds;
    parent[root] = -1;
    nds.push(root);
    while (!nds.empty())
    {
        int x = nds.front(); nds.pop();
        for (int y: new_edges[x])
            if (parent[y] == -2)
            {
                parent[y] = x;
                int c = color[x];
                if (x == root)
                    c = y;
                color[y] = c;
                color_count[c]++;
                nds.push(y);
            }
    }

    int big_alrdy = -1;
    for (int i=0; i<n; i++)
        if (color_count[i] >= a)
            big_alrdy = i;
    
    if (big_alrdy != -1)
    {
        //cout << "malaki na po ako\n";
        vector<int> colorby2(n, 5);
        for (int i=0; i<n; i++)
            if (color[i] == big_alrdy)
                colorby2[i] = 4;
        return find_split_given_two_colors(n, a, b, c, edges, colorby2);
    }

    //cout << "reached here\n";
    
    vector<vector<int>> color_edge(n);
    for (int i=0; i<e; i++)
    {
        int ppp = color[p[i]];
        int qqq = color[q[i]];
        if (ppp!=qqq)
        {
            color_edge[ppp].push_back(qqq);
            color_edge[qqq].push_back(ppp);
        }
    }

    int total = color_count[color[root_a]];
    queue<int> colors_in;
    vector<bool> colorby4(n, 0);
    colors_in.push(color[root_a]);
    colorby4[color[root_a]] = 1;
    while (total < a)
    {
        int x = colors_in.front(); colors_in.pop();
        for (int c: color_edge[x])
            if ((colorby4[c]==0) and (total < a))
            {
                total += color_count[c];
                colors_in.push(c);
                colorby4[c] = 1;
            }
    }

    vector<int> colorby2(n, 5);
    for (int i=0; i<n; i++)
        if (colorby4[color[i]])
            colorby2[i] = 4;
    return find_split_given_two_colors(n, a, b, c, edges, colorby2);
    
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    //0 2 3 1
    //112223
    //223331
	vector<int> rewrite = {0, 1, 2, 3};
    for (int i=0; i<5; i++)
    {
        if (a>b)
        {
            swap(a, b);
            swap(rewrite[1], rewrite[2]);
        }
        if (b>c)
        {
            swap(b, c);
            swap(rewrite[2], rewrite[3]);
        }
    }

    vector<int> ans = find_split_wlog(n, a, b, c, p, q);
    for (int i=0; i<n; i++)
    {
        ans[i] = rewrite[ans[i]];
        /*if (ans[i]==rewrite[0])
            ans[i] = 0;
        else if (ans[i]==rewrite[1])
            ans[i] = 1;
        else if (ans[i]==rewrite[2])
            ans[i] = 2;
        else if (ans[i]==rewrite[3])
            ans[i] = 3;*/
    }
    return ans;
}

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