Submission #1035674

#	Time	Username	Problem	Language	Result	Execution time	Memory
1035674		Sam_a17	Split the Attractions (IOI19_split)	C++17	0 / 100	516 ms	1048576 KiB

split

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

// -------------------- Typedefs -------------------- //

typedef long long ll;
typedef unsigned long long ull;
typedef double db;
typedef long double ld;

// -------------------- Defines -------------------- //

#define pr  pair
#define mpr make_pair
#define ff  first
#define ss  second

#define mset   multiset
#define mmap   multimap
#define uset   unordered_set
#define umap   unordered_map
#define umset  unordered_multiset
#define ummap  unordered_multimap
#define pqueue priority_queue

#define sz(x)  (int((x).size()))
#define len(x) (int((x).length()))
#define all(x) (x).begin(), (x).end()
#define clr(x) (x).clear()

#define ft   front
#define bk   back
#define pf   push_front
#define pb   push_back
#define popf pop_front
#define popb pop_back

#define lb lower_bound
#define ub upper_bound
#define bs binary_search

// -------------------- Constants -------------------- //

const int MAX = int(1e9 + 5);
const ll MAXL = ll(1e18 + 5);
const ll MOD = ll(1e9 + 7);
const ll MOD2 = ll(998244353);

// -------------------- Solution -------------------- //

const int N = 300005;
vector<int> g[N];
int sz[N];
int ans[N];
int n, m;

int timer, tin[N], tout[N];

set<pr<int, int>> e1[N];
vector<int> e2[N];

void dfs(int u, int p = 0)
{
    int i, j;
    int v;

    tin[u] = ++timer;

    for (i = 0; i < sz(g[u]); i++)
    {
        v = g[u][i];

        if (v == p)
            continue;

        dfs(v, u);

        sz[u] += sz[v];
    }

    sz[u]++;

    tout[u] = ++timer;

    return;
}

void dfs0(int u, int p = 0)
{
    int i, j;
    int v;

    e1[sz[u]].insert(mpr(tin[u], u));

    for (i = 0; i < sz(g[u]); i++)
    {
        v = g[u][i];
        if (v == p)
            continue;
        dfs0(v, u);
    }

    return;
}

pr<pr<int, int>, int> dfs1(int u, int p, int& a, int& b, int& c)
{
    int i, j;
    int v;

    if (e1[sz[u]].find(mpr(tin[u], u)) != e1[sz[u]].end())
    {
        e1[sz[u]].erase(mpr(tin[u], u));
    }
    e2[sz[u]].pb(u);

    if (sz[u] == a)
    {
        if (!e1[b].empty())
        {
            auto it = e1[b].begin();
            if ((*it).ff < tin[u])
                return mpr(mpr(u, (*it).ss), 1);

            it = e1[b].end();
            it--;
            if ((*it).ff > tout[u])
                return mpr(mpr(u, (*it).ss), 1);
        }
        if (!e2[a + b].empty())
            return mpr(mpr(u, e2[a + b].ft()), 2);
    }

    for (i = 0; i < sz(g[u]); i++)
    {
        v = g[u][i];

        if (v == p)
            continue;

        pr<pr<int, int>, int> res = dfs1(v, u, a, b, c);
        if (res.ss)
            return res;
    }

    e1[sz[u]].insert(mpr(tin[u], u));
    e2[sz[u]].popb();

    return mpr(mpr(0, 0), 0);
}

void check(int a, int b, int c)
{
    int i, j;

    for (i = 1; i <= n; i++)
    {
        clr(e1[i]);
        clr(e2[i]);
    }

    dfs0(1);

    pr<pr<int, int>, int> res = dfs1(1, 0, a, b, c);

    if (res.ss == 0)
        return;

    // cout << res.ff.ff << ' ' << res.ff.ss << ' ' << res.ss << ' ' << a << ' ' << b << ' ' << c << "...\n";

    for (i = 1; i <= n; i++)
        ans[i] = 1;

    if (res.ss == 1)
    {
        for (i = 1; i <= n; i++)
        {
            if (tin[res.ff.ff] <= tin[i] && tout[i] <= tout[res.ff.ff])
                ans[i] = 2;
            if (tin[res.ff.ss] <= tin[i] && tout[i] <= tout[res.ff.ss])
                ans[i] = 3;
        }
    }
    else
    {
        for (i = 1; i <= n; i++)
        {
            if (tin[res.ff.ss] <= tin[i] && tout[i] <= tout[res.ff.ss])
                ans[i] = 2;
            if (tin[res.ff.ff] <= tin[i] && tout[i] <= tout[res.ff.ff])
                ans[i] = 3;
        }
    }

    return;
}

vector<int> find_split(int n0, int a, int b, int c, vector<int> p0, vector<int> q0)
{
    int i, j;

    n = n0;
    m = sz(p0);

    for (i = 0; i < m; i++)
    {
        g[p0[i] + 1].pb(q0[i] + 1);
        g[q0[i] + 1].pb(p0[i] + 1);
    }

    dfs(1);

    // for (i = 1;i <= n;i++) cout << sz[i] << ' ';
    // cout << "aa\n";

    check(a, b, c);
    // if (!ans[1])
    //     check(a, c, b);
    // if (!ans[1])
    //     check(b, a, c);
    // if (!ans[1])
    //     check(b, c, a);
    // if (!ans[1])
    //     check(c, a, b);
    // if (!ans[1])
    //     check(c, b, a);

    vector<int> res;
    for (i = 1; i <= n; i++)
        res.pb(ans[i]);

    if (!ans[1])
        return res;

    int cnt[4] = {0, 0, 0, 0}, cc[4];
    vector<int> tor = {0, 1, 2, 3};

    for (i = 0; i < n; i++)
        cnt[res[i]]++;

    // cout << cnt[1] << ' ' << cnt[2] << ' ' << cnt[3] << ".\n";

    do
    {
        // cout << tor[1] << ' ' << tor[2] << ' ' << tor[3] << "\n";
        if (cnt[tor[1]] == a && cnt[tor[2]] == b && cnt[tor[3]] == c)
            break;
    } while (next_permutation(tor.begin() + 1, tor.end()));

    for (i = 1; i <= 3; i++)
        cc[tor[i]] = i;

    for (i = 0; i < n; i++)
        res[i] = cc[res[i]];

    return res;
}

/*
    # # # #   # # # #   # # # #   #       #    #       #     #
       #      #         #     #    #     #    # #      #   #
       #      # # # #   #     #     #   #    #   #     # #
       #            #   #     #      # #    # # # #    #   #
       #      # # # #   # # # #       #    #       #   #     #
*/

Compilation message (stderr)

split.cpp: In function 'void dfs(int, int)':
split.cpp:67:12: warning: unused variable 'j' [-Wunused-variable]
   67 |     int i, j;
      |            ^
split.cpp: In function 'void dfs0(int, int)':
split.cpp:93:12: warning: unused variable 'j' [-Wunused-variable]
   93 |     int i, j;
      |            ^
split.cpp: In function 'std::pair<std::pair<int, int>, int> dfs1(int, int, int&, int&, int&)':
split.cpp:111:12: warning: unused variable 'j' [-Wunused-variable]
  111 |     int i, j;
      |            ^
split.cpp: In function 'void check(int, int, int)':
split.cpp:157:12: warning: unused variable 'j' [-Wunused-variable]
  157 |     int i, j;
      |            ^
split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:203:12: warning: unused variable 'j' [-Wunused-variable]
  203 |     int i, j;
      |            ^

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