Submission #589432

#	Time	Username	Problem	Language	Result	Execution time	Memory
589432		PiejanVDC	Split the Attractions (IOI19_split)	C++17	29 / 100	465 ms	1048576 KiB

split

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

int A,B;
int N;

vector<bool>vis;
int need;

vector<int>adj[100005];

int C;
int c_a, c_b, c_c;

vector<int>answer;

void D(int u) {
    if(!need)
        return;
    vis[u] = 1;
    answer[u] = C;
    need--;
    for(auto z : adj[u]) if(!vis[z]) {
        D(z);
    }
}

void color(int s1, int n1, int s2, int up) {
    vis.clear();
    vis.resize(N, 0);
    if(n1 != -1)
        vis[n1] = 1;
    if(up != -1)
        vis[up] = 1;
    need = B;
    C = c_b;
    D(s1);
    C = c_a, need = A;
    D(s2);
    for(auto &z : answer)
        if(!z)
            z = c_c;
}

int dfs(int u, int e = -1) {

    //cout << u << ' ';

    int sz = 0;
    vector<pair<int,int>>b;
    int total = 1;
    for(auto z : adj[u]) if(z != e) {
        b.emplace_back(dfs(z, u), z);
        if(b.back().first == -1)
            return -1;
        total += b.back().first;
    }
    sort(b.begin(), b.end());

    // binary search for A, sum the rest for B

    int l = 0, r = (int)b.size() - 1;
    int p = -1;

    while(l <= r) {
        int mid = (l + r)/2;
        if(b[mid].first >= A)
            p = mid, r = mid-1;
        else
            l = mid+1;
    }

    if(p != -1) {
        if(total - b[p].first >= B) {
            color(u, b[p].second, b[p].second, e);
            return -1;
        }
    }

    // check if total - sum here >= B and sum here >= A

    if(total >= B && N - total >= A) {
        color(u, -1, e, e);
        return -1;
    }

    return total;
}

void solve(vector<pair<int,int>>&x) {
    A = x[0].first, B = x[1].first;
    c_a = x[0].second, c_b = x[1].second, c_c = x[2].second;
    if(dfs(0) == -1)
        sort(x.rbegin(), x.rend());
}

vector<int>find_split(int n, int a, int b, int c, vector<int>p, vector<int>q) {
    N = n;
    answer.resize(N, 0);
    vector<pair<int,int>>x = {{a,1}, {b,2}, {c,3}};
    sort(x.begin(), x.end());

    for(int i = 0 ; i < n - 1 ; i++)
        adj[p[i]].push_back(q[i]), adj[q[i]].push_back(p[i]);

    do {
        solve(x);
    } while(next_permutation(x.begin(), x.end()));
    return answer;
}

Compilation message (stderr)

split.cpp: In function 'int dfs(int, int)':
split.cpp:50:9: warning: unused variable 'sz' [-Wunused-variable]
   50 |     int sz = 0;
      |         ^~

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