Submission #1070564

#	Time	Username	Problem	Language	Result	Execution time	Memory
1070564		j_vdd16	Split the Attractions (IOI19_split)	C++17	40 / 100	78 ms	24664 KiB

split

#include "split.h"

#include <algorithm>
#include <bitset>
#include <cstdint>
#include <cstring>
#include <iostream>
#include <limits.h>
#include <math.h>
#include <map>
#include <numeric>
#include <queue>
#include <set>
#include <stack>
#include <string>
#include <vector>

//#define int long long
#define loop(X, N) for(int X = 0; X < (N); X++)
#define all(V) V.begin(), V.end()
#define rall(V) V.rbegin(), V.rend()

using namespace std;

typedef vector<int> vi;
typedef vector<vi> vvi;
typedef pair<int, int> ii;
typedef vector<ii> vii;
typedef vector<vector<ii>> vvii;
typedef vector<bool> vb;
typedef vector<vector<bool>> vvb;

int n, m;
int a, b, c;
vvi adj;
vvi adjTree;
vi res;

vb vis;
vi subSize;
ii marked1, marked2;
vector<ii> types;
int dfs2(int node, int parent) {
    vis[node] = true;

    int maxSize = 0;
    int childCount = 0;
    for (int child : adj[node]) {
        if (vis[child]) continue;

        if (childCount >= 2) continue;
        childCount++;

        adjTree[node].push_back(child);
        adjTree[child].push_back(node);

        int childSz = dfs2(child, node);
        maxSize = max(maxSize, childSz);

        subSize[node] += childSz;
    }
    subSize[node]++;

    if (maxSize < types[2].first && subSize[node] >= types[2].first && n - subSize[node] >= types[1].first) {
        marked2 = { node, parent };
    } 
    else if (maxSize < types[1].first && subSize[node] >= types[1].first && n - subSize[node] >= types[2].first) {
        marked1 = { node, parent };
    }

    return subSize[node];
}
int dfs(int node, int parent) {
    vis[node] = true;

    int maxSize = 0;
    for (int child : adj[node]) {
        if (vis[child]) continue;

        adjTree[node].push_back(child);
        adjTree[child].push_back(node);

        int childSz = dfs(child, node);
        maxSize = max(maxSize, childSz);

        subSize[node] += childSz;
    }
    subSize[node]++;

    if (maxSize < types[2].first && subSize[node] >= types[2].first && n - subSize[node] >= types[1].first) {
        marked2 = { node, parent };
    } 
    else if (maxSize < types[1].first && subSize[node] >= types[1].first && n - subSize[node] >= types[2].first) {
        marked1 = { node, parent };
    }

    return subSize[node];
}

void partition(int node, int type, int parent, ii marked) {
    auto& [nodesLeft, value] = types[type];
    if (nodesLeft == 0) {
        type = 3;
        res[node] = types[type].second;
    }
    else {
        nodesLeft--;
        res[node] = value;
    }

    for (int child : adjTree[node]) {
        if (child == parent) continue;

        if (node == marked.first) {
            if (child == marked.second) {
                partition(child, 2, node, marked);
            }
            else {
                partition(child, 1, node, marked);
            }
        }
        else {
            partition(child, type, node, marked);
        }
    }
}

vector<int> find_split(int _n, int _a, int _b, int _c, vector<int> p, vector<int> q) {
    n = _n, m = (int)p.size();
    adj = vvi(n);
    adjTree = vvi(n);
    a = _a, b = _b, c = _c;

    loop(i, m) {
        adj[p[i]].push_back(q[i]);
        adj[q[i]].push_back(p[i]);
    }

    vis = vb(n);
    subSize = vi(n);
    marked1 = { -1, -1 }, marked2 = { -1, -1 };
    types = {{0, 0}, {a, 1}, {b, 2}, {c, 3}};
    sort(all(types));
    dfs2(0, -1);

    int visCount = 0;
    loop(i, n) {
        visCount += vis[i];
    }

    res = vi(n);
    if (visCount != n || marked1 == ii{-1, -1} && marked2 == ii{-1, -1}) {
        vis = vb(n);
        subSize = vi(n);
        marked1 = { -1, -1 }, marked2 = { -1, -1 };
        types = {{0, 0}, {a, 1}, {b, 2}, {c, 3}};
        adjTree = vvi(n);
        sort(all(types));
        dfs(0, -1);

        if (marked1 == ii{-1, -1} && marked2 == ii{-1, -1}) {
            return res;
        }
        ii marked = marked1;
        if (marked1 == ii{-1, -1}) {
            marked = marked2;
            swap(types[1], types[2]);
        }

        partition(marked.first, 1, -1, marked);
        return res;
    }
    ii marked = marked1;
    if (marked1 == ii{-1, -1}) {
        marked = marked2;
        swap(types[1], types[2]);
    }

    partition(marked.first, 1, -1, marked);
	return res;
}

Compilation message (stderr)

split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:152:48: warning: suggest parentheses around '&&' within '||' [-Wparentheses]
  152 |     if (visCount != n || marked1 == ii{-1, -1} && marked2 == ii{-1, -1}) {
      |                          ~~~~~~~~~~~~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~~~~

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