Submission #1229877

#TimeUsernameProblemLanguageResultExecution timeMemory
1229877tapilyocaSplit the Attractions (IOI19_split)C++20
22 / 100
590 ms1114112 KiB
#include "split.h"
#include <bits/stdc++.h>
using namespace std;

template<typename T>
using vec = vector<T>;
using ll = long long;
using vll = vec<ll>;
using vvll = vec<vll>;
using pll = pair<ll,ll>;
#define pb push_back
#define dbg(x) cerr << #x << ": " << x << endl


vll subtree_size;
vec<int> color;
vec<pll> days;
vll dayOrder(3);
vvll adj;
pll good_edge = {-1,-1};
ll N;

void dfs(ll u, ll p) {
    // cerr << "Call " << u << " " << p << endl;
    subtree_size[u] = 1;
    for(ll &v : adj[u]) {
        if(v == p) continue;
        dfs(v,u);
        subtree_size[u] += subtree_size[v];
    }

    // can we cut (u,p) ?
    if(subtree_size[u] >= days[0].first and N - subtree_size[u] >= days[1].first) {
        // cerr << "Yay, " << u << " " << p << endl;
        good_edge = {u,p};
    }
    else if(subtree_size[u] >= days[1].first and N - subtree_size[u] >= days[0].first) {
        good_edge = {p,u};
    }
}

void dfs_color(ll u, ll p, ll c) {
    if(days[dayOrder[c-1]].first == 0) return;
    
    color[u] = c;
    days[dayOrder[c-1]].first--;
    
    for(ll &v : adj[u]) {
        if(v == p) continue;
        dfs_color(v,u,c);
    }
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
	N = n;
    ll m = p.size();
    adj.clear();
    adj.resize(n);

    for(int i = 0; i < m; i++) {
        adj[p[i]].pb(q[i]);
        adj[q[i]].pb(p[i]);
        // cerr << "Edge " << p[i] << " " << q[i] << endl;
    }

    days = {{a,1}, {b,2}, {c,3}};
    sort(days.begin(),days.end());
    for(int i = 0; i < 3; i++) dayOrder[days[i].second - 1] = i;

    subtree_size.assign(n,0);
    color.assign(n,-1);
    dfs(0,-1);

    if(good_edge.first == -1) {
        vec<int> out(n,0);
        return out;
    }

    dfs_color(good_edge.first, good_edge.second, days[0].second);


    // cerr << "Days: ";
    // for(pll &x :days) cerr << "(" << x.first << ", " << x.second << "), "; 
    // cerr << endl;
    // cerr << "Color so far: " << endl;
    // for(int &x : color) cerr << x << " ";
    // cerr << endl;
    
    dfs_color(good_edge.second, good_edge.first, days[1].second);

    for(int &x : color) x = (x != -1 ? x : days[2].second);
    return color;
}
#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...