Submission #832969

#	Time	Username	Problem	Language	Result	Execution time	Memory
832969		HorizonWest	Split the Attractions (IOI19_split)	C++17	40 / 100	392 ms	1048576 KiB

split

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

#define fs first
#define sd second

vector <vector<int>> v;
vector <int> g, p1, dp, ans, pass;
vector <pair<int, int>> gx;
int N; bool can = false;

int last(int node)
{
    p1[node] = 1;
    for(auto& u : v[node])
        if(!p1[node]) return last(u);
    return node;
}

void dfs(int node, int& cont, int& group)
{
    ans[node] = group; cont++;
    if(cont == g[group])
    {
        cont = 0;
        group = max(1, (group + 1) % 4);
    }
    pass[node] = 1;
    for(auto& u : v[node])
    {
        if(!pass[u])
        {
            dfs(u, cont, group);
        }
    }

    return;
}

void dfs2(int node, int parent)
{
    for(auto& u : v[node]) if(u != parent)
    {
        dfs2(u, node);
        dp[node] += dp[u];
    }
}

void dfs3(int node, int parent, int& cont, int& group)
{
    if(cont == gx[group].fs) return;
    cont++; ans[node] = gx[group].sd;
    for(auto& u : v[node]) if(u != parent)
    {
        dfs3(u, node, cont, group);
    }
}

void dfs4(int node, int parent)
{
    if(pass[node]) return;
    pass[node] = 1;
    for(auto& u : v[node]) if(u != parent)
    {
        if(dp[u] >= gx[0].fs && N - dp[u] >= gx[1].fs)
        { can = true;
            int cont = 0, group = 0;
            dfs3(u, node, cont, group);
            cont = 1; group = 1; ans[node] = gx[1].sd;
            for(auto& j : v[node]) if(j != u)
            {
                dfs3(j, node, cont, group);
            }
            return;
        }

        if(dp[u] >= gx[1].fs && N - dp[u] >= gx[0].fs)
        { can = true;
            int cont = 0, group = 1;
            dfs3(u, node, cont, group);
            cont = 1; group = 0; ans[node] = gx[0].sd;
            for(auto& j : v[node]) if(j != u)
            {
                dfs3(j, node, cont, group);
            }
            return;
        }

        dfs4(u, node);
    }
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q)
{
    int m = p.size(); ans.assign(n, 3); pass.assign(n, 0); N = n;

    v.assign(n + 1, vector <int> ());
    vector <bool> subtask(5, 1);

    for(int i = 0; i < m; i++)
    {
        v[p[i]].push_back(q[i]);
        v[q[i]].push_back(p[i]);

        if(v[p[i]].size() > 2) subtask[0] = false;
        if(v[q[i]].size() > 2) subtask[0] = false;
    }

    if(a != 1) subtask[1] = false;

    if(subtask[0] || subtask[1])
    {
        p1.assign(n, 0); g.assign(4, 0);
        g[1] = a; g[2] = b; g[3] = c;
        int cont = 0, group = 2;
        dfs(last(0), cont, group);
        return ans;
    }

    if(subtask[2])
    {
        gx = {{ a, 1 }, { b, 2 }, { c, 3 }};
        ans.clear(); ans.assign(n, gx[2].sd);
        dp.assign(n, 1); dfs2(0, -1);
        sort(gx.begin(), gx.end());
        dfs4(0, -1);

        if(!can == 1)
        {
            ans.clear();
            ans.assign(n, 0);
        }

        return ans;
    }

	return ans;
}

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