Submission #253748

#TimeUsernameProblemLanguageResultExecution timeMemory
253748TriPhanSplit the Attractions (IOI19_split)C++14
100 / 100
193 ms65272 KiB
#include <bits/stdc++.h>
#include <split.h>
#define TASK "SPLIT"
#define pb push_back
#define ii pair<int,int>
#define iii pair<ii,int>
#define ll long long
#define F first
#define S second
#define FOR(i, a, b) for (int i=a; i<=b; i++)
#define FOr(i, a, b) for (int i=a; i<b ; i++)
#define FOD(i, a, b) for (int i=a; i>=b; i--)
#define FOd(i, a, b) for (int i=a; i>b ; i--)
#define vi vector<int>
using namespace std;

const int N=(int)1e6+7;
const int MOD=(int)1e9+7;
const ll INF=(ll)1e18+7;
vector<int> a[N],e[N];
ii A[3];
int visit[N],par[N],sz[N],ans[N],del[N],lift[N];
int m,n,T,root;
void dfs(int u) {
    sz[u]=1;
    for (int v:a[u]) {
        if (v==par[u])
            continue;
        if (visit[v]) {
            if (lift[v]==0 || visit[lift[u]]>visit[v])
                lift[u]=v;
            continue;
        }
        visit[v]=visit[u]+1;
        par[v]=u;
        dfs(v);
        sz[u]+=sz[v];
    }
    if (sz[u] >= A[0].F && root==0)
        root=u;
}

void DFS(int u,int p) {
    if (A[T].F == 0) return;

        A[T].F--;
        ans[u]=A[T].S;

    visit[u]=1;
    for (int v:a[u])
        if (v!=p && visit[v]==0)
            DFS(v,u);
}
void DFS2(int u,int p) {
    if (A[T].F == 0) return;
    A[T].F--;
    ans[u]=A[T].S;

    for (int v:a[u])
        if (!ans[v])
            DFS2(v,u);
}
vi find_split(int n, int aa, int b, int c, vi p, vi q) {

    m=p.size();
    A[0].F=aa;
    A[1].F=b;
    A[2].F=c;
    A[0].S=1;
    A[1].S=2;
    A[2].S=3;
    sort(A,A+3);
    FOR(i,0,m-1) {
        int u,v;
        u=p[i];
        v=q[i];
        ++u; ++v;
        a[u].pb(v);
        a[v].pb(u);
    }
    visit[1]=1;
    dfs(1);
    vector<int> res;
    if (par[root]==0) {
        FOR(i,1,n)
            res.push_back(0);
        return res;
    }

    queue<ii> Q;
    for (int v:a[root])
        if (par[v]==root)
            Q.push({v,v});
    while (n-sz[root] < A[0].F && Q.size()) {
        ii z=Q.front();
        int v=z.F;
        int r=z.S;
        Q.pop();
        if (del[r]) continue;
        if (lift[v] && visit[lift[v]] < visit[root]) {
            par[v]=lift[v];
            par[r]=0;
            sz[root]-=sz[r];
            del[r]=1;
            continue;
        }
        for (int u:a[v])
            if (par[u]==v)
              Q.push({u,r});

    }
    if (n-sz[root] < A[0].F) {
        FOR(i,1,n)
            res.push_back(0);
        return res;
    }
    FOR(i,1,n) {
        e[i]=a[i];
        a[i].clear();
    }
    FOR(i,1,n)
        if (i!=root && par[i]!=0) {
            a[par[i]].push_back(i);
            a[i].push_back(par[i]);
        }
    FOR(i,1,n)
        visit[i]=0;

    if (sz[root]>=A[1].F) {
        T=1;
        DFS(root,0);
        FOR(i,1,n)
        a[i]=e[i];
        T=0;
        DFS2(1,0);
    }
    else {
        T=0;
        DFS(root,0);
        FOR(i,1,n)
        a[i]=e[i];
        T=1;
        DFS2(1,0);
    }
    FOR(i,1,n)
        if (!ans[i])
            ans[i]=A[2].S;
    FOR(i,1,n)
        res.push_back(ans[i]);
    return res;
}
#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...