Submission #1203827

#TimeUsernameProblemLanguageResultExecution timeMemory
1203827AvianshSplit the Attractions (IOI19_split)C++20
22 / 100
60 ms14284 KiB
#include "split.h"
#include <bits/stdc++.h>

using namespace std;

int rt;

void dfs(int st, vector<int>g[], int sz[], int a, int b, int c, int p){
    for(int i : g[st]){
        if(sz[i]>=a&&sz[st]-sz[i]>=min(b,c)){
            rt=st;
        }
        if(sz[i]>=b&&sz[st]-sz[i]>=min(a,c)){
            rt=st;
        }
        if(sz[i]>=c&&sz[st]-sz[i]>=min(b,a)){
            rt=st;
        }
        if(i==p)
            continue;
        int n = sz[st];
        int o = sz[i];
        sz[st]-=sz[i];
        sz[i]=n;
        dfs(i,g,sz,a,b,c,st);
        sz[i]=o;
        sz[st]=n;
    }
}

void dfs_sz(int st, vector<int>g[], int sz[], int p){
    sz[st]=1;
    for(int i : g[st]){
        if(i==p)
            continue;
        dfs_sz(i,g,sz,st);
        sz[st]+=sz[i];
    }
}

vector<int>order;

void dfso(int st, vector<int>g[], bool vis[]){
    order.push_back(st);
    vis[st]=1;
    for(int i : g[st]){
        if(vis[i])
            continue;
        dfso(i,g,vis);
    }
}

vector<int> find_split(int n, int a, int b, int c, vector<int> p, vector<int> q) {
    vector<int>ans(n,0);
    vector<int>g[n];
    int m = p.size();
    for(int i = 0;i<m;i++){
        g[p[i]].push_back(q[i]);
        g[q[i]].push_back(p[i]);
    }
    assert(m==n-1);
    rt=-1;
    int sz[n];
    dfs_sz(0,g,sz,-1);
    dfs(0,g,sz,a,b,c,-1);
    if(rt==-1){
        return ans;
    }
    dfs_sz(rt,g,sz,-1);
    //sz reset according to rt
    for(int i : g[rt]){
        int st = rt;
        if(sz[i]>=a&&sz[st]-sz[i]>=min(b,c)){
            //must set all to a
            bool vis[n];
            fill(vis,vis+n,0);
            vis[rt]=1;
            order.clear();
            dfso(i,g,vis);
            //order has all elements now
            //must set a of them to 1
            for(int i = 0;i<a;i++){
                ans[order[i]]=1;
            }
            //now must set other parts to val according to min (b,c)
            order.clear();
            vis[rt]=0;
            dfso(rt,g,vis);
            if(b<c){
                for(int i = 0;i<b;i++){
                    ans[order[i]]=2;
                }
                for(int i = 0;i<n;i++){
                    if(ans[i]==0){
                        ans[i]=3;
                    }
                }
            }
            else{
                for(int i = 0;i<c;i++){
                    ans[order[i]]=3;
                }
                for(int i = 0;i<n;i++){
                    if(ans[i]==0){
                        ans[i]=2;
                    }
                }
            }
            return ans;
        }
        if(sz[i]>=b&&sz[st]-sz[i]>=min(a,c)){
            //must set all to b
            bool vis[n];
            fill(vis,vis+n,0);
            vis[rt]=1;
            order.clear();
            dfso(i,g,vis);
            //order has all elements now
            //must set a of them to 1
            for(int i = 0;i<b;i++){
                ans[order[i]]=2;
            }
            //now must set other parts to val according to min (b,c)
            order.clear();
            vis[rt]=0;
            dfso(rt,g,vis);
            if(a<c){
                for(int i = 0;i<a;i++){
                    ans[order[i]]=1;
                }
                for(int i = 0;i<n;i++){
                    if(ans[i]==0){
                        ans[i]=3;
                    }
                }
            }
            else{
                for(int i = 0;i<c;i++){
                    ans[order[i]]=3;
                }
                for(int i = 0;i<n;i++){
                    if(ans[i]==0){
                        ans[i]=1;
                    }
                }
            }
            return ans;
        }
        if(sz[i]>=c&&sz[st]-sz[i]>=min(b,a)){

            //must set all to c
            bool vis[n];
            fill(vis,vis+n,0);
            vis[rt]=1;
            order.clear();
            dfso(i,g,vis);
            //order has all elements now
            //must set a of them to 1
            for(int i = 0;i<c;i++){
                ans[order[i]]=3;
            }
            //now must set other parts to val according to min (b,c)
            order.clear();
            vis[rt]=0;
            dfso(rt,g,vis);
            if(b<a){
                for(int i = 0;i<b;i++){
                    ans[order[i]]=2;
                }
                for(int i = 0;i<n;i++){
                    if(ans[i]==0){
                        ans[i]=1;
                    }
                }
            }
            else{
                for(int i = 0;i<a;i++){
                    ans[order[i]]=1;
                }
                for(int i = 0;i<n;i++){
                    if(ans[i]==0){
                        ans[i]=2;
                    }
                }
            }
            return ans;
        }
    }
    return ans;
}
#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...