Submission #298849

#TimeUsernameProblemLanguageResultExecution timeMemory
298849Aldas25Split the Attractions (IOI19_split)C++14
40 / 100
138 ms15992 KiB
#include "split.h"
#include <bits/stdc++.h>

using namespace std;

#define FAST_IO ios_base::sync_with_stdio(0); cin.tie(nullptr)
#define FOR(i, a, b) for (int i = (a); i <= (b); i++)
#define REP(n) FOR(O, 1, (n))
#define f first
#define s second
#define pb push_back
typedef long long ll;
typedef pair<int, int> pii;
typedef vector<int> vi;
typedef vector<pii> vii;
typedef vector<ll> vl;

const int MAXN = 100100;

vi adj[MAXN];
vi seq;
bool ch[MAXN];
int col[3];
int n;

void dfs (int v) {
    seq.pb(v);
    ch[v] = true;
    for (int u : adj[v]) {
        if (!ch[u])
            dfs(u);
    }
}

int sz[MAXN];
bool found = false;
pii paint1, paint2;

void dfsTree (int v,int p) {
    if (found) return;
    sz[v] ++;
    for (int u : adj[v]) {
        if (u == p) continue;
        dfsTree(u,v);
        sz[v] += sz[u];
    }
    if (found) return;

   // cout << " v= " << v << " sz=" << sz[v] << " n - sz = " << n-sz[v] << endl;

    FOR(i, 0, 2) FOR(j, 0, 2) {
        if (i == j) continue;
        if (sz[v] >= col[i] && (n-sz[v]) >= col[j]) {
            found = true;
            paint1 = {v, i};
            paint2 = {p, j};
            break;
        }
    }

}

int painted[MAXN];

void paint (int v, int p, int c) {
    if (col[c] > 0) {
        painted[v] = c+1;
        col[c]--;
    }
    for (int u : adj[v]) {
        if (u == p) continue;
        paint(u, v, c);
    }
}

vector<int> find_split(int N, int a, int b, int c, vector<int> p, vector<int> q) {
	n = N;
	vector<int> res(n, 0);

	col[0] = a, col[1] = b, col[2] = c;

    int m = (int)p.size();
    FOR(i, 0, m-1) {
        int u = p[i], v = q[i];
        adj[u].pb(v);
        adj[v].pb(u);
    }

    if (m == n-1) {
        dfsTree (0, -1);
        if (!found) return res;

        FOR(i, 0, 2) {
            if (i == paint1.s || i == paint2.s) continue;
            FOR(j, 0, n-1) res[j] = i+1;
        }

        //cout << " paint1: " << paint1.f << "  col:" << paint1.s << endl;
        //cout << " paint2: " << paint2.f << "  col:" << paint2.s << endl;


        paint(paint1.f, paint2.f, paint1.s);
        paint(paint2.f, paint1.f, paint2.s);
        FOR(i, 0, n-1) if (painted[i] > 0) res[i] = painted[i];
        return res;
    }

    int st = 0;
    FOR(i, 0, n-1) if ((int)adj[i].size() == 1) st = i;
    dfs (st);

    FOR(i, 0, a-1) res[seq[i]] = 1;
    FOR(i, a, a+b-1) res[seq[i]] = 2;
    FOR(i, a+b, a+b+c-1) res[seq[i]] = 3;

	return res;
}

/*
9 10
4 2 3
0 1
0 2
0 3
0 4
0 6
0 8
1 7
3 7
4 5
5 6
ans:
1 1 3 1 2 2 3 1 3

6 5
2 2 2
0 1
0 2
0 3
0 4
0 5
ans:
0 0 0 0 0 0


9 8
4 2 3
0 1
0 2
0 3
0 4
0 8
1 7
4 5
5 6
ans:
1 1 3 1 2 2 3 1 3

*/

#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...