Submission #265142

#	Time	Username	Problem	Language	Result	Execution time	Memory
265142		hamerin	Split the Attractions (IOI19_split)	C++17	0 / 100	190 ms	36480 KiB

split

#include "split.h"

#include <bits/stdc++.h>

#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")

using namespace std;

using i64 = long long;
using d64 = long double;
using pi = pair<int, int>;
using pli = pair<i64, i64>;
using ti = tuple<int, int, int>;
using tli = tuple<i64, i64, i64>;

#define iterall(cont) cont.begin(), cont.end()
#define prec(n) setprecision(n) << fixed

class disjointSet {
   public:
    vector<int> p;

    disjointSet() = default;
    explicit disjointSet(int N) {
        p.clear();
        p.resize(N);
		for(int i=0;i<N;i++) p[i] = i;
    }

    int find(int u) { return p[u] = (p[u] == u ? u : find(p[u])); }

    void mer(int u, int v) { p[find(v)] = find(u); }

    bool sset(int u, int v) { return find(u) == find(v); }
};

namespace GRAPH {

int V, E, cen;
vector<int> sut, sz, par;
vector<pi> edges;
vector<vector<int>> adj, tr;
vector<bool> vst;
disjointSet dS;

void init(int _V) {
    V = _V;
    sut.resize(V), sz.resize(V), par.resize(V);
    adj.resize(V), tr.resize(V);
    dS = disjointSet(V);
}

void emplace_edge(int u, int v) {
    ++E;
    edges.emplace_back(u, v);
    adj[u].emplace_back(v);
    adj[v].emplace_back(u);
}

// dfs ordering && parent
void dfs1(int h, bool ini = false) {
    if (ini) {
        vst.clear();
        vst.resize(V);
        vst[h] = true;
        par[h] = -1;
    }

    for (auto t : adj[h]) {
        if (vst[t]) continue;

        vst[t] = true;
        tr[h].emplace_back(t);
        par[t] = h;
        dfs1(t);
    }
}

// subtree size
void dfs2(int h) {
    sz[h] = 1;
    for (auto t : tr[h]) {
        dfs2(t);
        sz[h] += sz[t];
    }
}

// find centroid
int dfs3(int h) {
    for (auto t : tr[h])
        if (sz[t] > (V >> 1)) return dfs3(t);
    return h;
}

bool trAdj(int u, int v) { return par[u] == v || par[v] == u; }

// centroid - dfs coloring
void dfs4(int h, int p) {
    for (auto t : adj[h]) {
        if (!trAdj(h, t)) continue;
        if (t == p) continue;

        if (h != cen) dS.mer(h, t);
		dfs4(t, h);
    }
}

void dfs() {
    dfs1(0, true);
    dfs2(0);
    cen = dfs3(0);
    dfs4(cen, -1);

    for (int i = 0; i < V; i++) sut[i] = dS.find(i);
}

void con(int h, int p, vector<int> &r, int &n) {
    if (n == r.size()) return;
    r[n++] = h;
	
    for (auto t : adj[h]) {
        if (!trAdj(h, t)) continue;
        if (t == p) continue;

        con(t, h, r, n);
    }
}
}  // namespace GRAPH

vector<int> find_split(int n, int a, int b, int c, vector<int> p,
                       vector<int> q) {
    vector<pi> targetSize = {pi(a, 1), pi(b, 2), pi(c, 3)};
    sort(iterall(targetSize));

    const size_t M = p.size();
    GRAPH::init(n);
    for (int i = 0; i < M; i++) GRAPH::emplace_edge(p[i], q[i]);
    GRAPH::dfs();

    auto newVer = GRAPH::sut;
    sort(iterall(newVer));

    vector<int> vertexSize(n);
    for (auto el : newVer) vertexSize[el]++;

    newVer.erase(unique(iterall(newVer)), newVer.end());
    int W = newVer.size();

    vector<int> aConf, bConf;

    // a이상 서브트리 존재
    for (int i = 0; i < W; i++) {
		if (newVer[i] == GRAPH::cen) continue;
        if (vertexSize[i] >= targetSize[0].first) {
            aConf.emplace_back(i);
            break;
        }
    }

    // a이상 서브트리 존재X
    if (aConf.empty()) {
        auto dS = disjointSet(W);

        for (auto [u, v] : GRAPH::edges) {
            if (GRAPH::trAdj(u, v)) continue;

            auto cu = *lower_bound(iterall(newVer), GRAPH::sut[u]);
            auto cv = *lower_bound(iterall(newVer), GRAPH::sut[v]);
            dS.mer(cu, cv);
        }

        vector<int> dSroot(W);
        for (int i = 0; i < W; i++) {
			if (newVer[i] == GRAPH::cen) continue;
			dSroot[i] = dS.find(i);
		}

        vector<int> sz(W);
        for (int i = 0; i < W; i++) {
			if (newVer[i] == GRAPH::cen) continue;
			sz[dSroot[i]] += vertexSize[newVer[i]];
		}

        for (int i = 0; i < W; i++) {
            if (sz[i] >= targetSize[0].first) {
                int vs = 0;

                for (int j = 0; j < W; j++) {
                    if (dSroot[j] == i)
                        aConf.emplace_back(j), vs += vertexSize[newVer[j]];
                    if (vs >= targetSize[0].first) break;
                }

                break;
            }
        }
    }

    {
        vector<bool> mask(W);
		mask[GRAPH::cen] = true;
        for (auto el : aConf) mask[el] = true;
        for (int i = 0; i < W; i++)
            if (!mask[i]) bConf.emplace_back(i);
    }

    // Configure
    if (aConf.empty()) return vector<int>(n, 0);

    vector<int> av(n, -1), bv(targetSize[1].first);
    int an = 0, bn = 1;
    bv[0] = GRAPH::cen;

    for (auto el : aConf) GRAPH::con(el, GRAPH::cen, av, an);
    for (auto el : bConf) GRAPH::con(el, GRAPH::cen, bv, bn);

	while(av.back() == -1) av.pop_back();
	vector<int> avv;
	{
		queue<int> que;
		vector<bool> vst(n, true);
		for(auto el:av) vst[el] = false;
		
		vst[av[0]] = true;
		que.emplace(av[0]);

		while(!que.empty() && avv.size() < targetSize[0].first) {
			int cur = que.front();
			que.pop();

			avv.emplace_back(cur);
			for(auto there:GRAPH::adj[cur]) {
				if(vst[there]) continue;
				que.emplace(there);
				vst[there] = true;
			}
		}
	}

    vector<int> ret(n, targetSize[2].second);
    for (auto el : avv) ret[el] = targetSize[0].second;
    for (auto el : bv) ret[el] = targetSize[1].second;

    return ret;
}

Compilation message (stderr)

split.cpp: In function 'void GRAPH::con(int, int, std::vector<int>&, int&)':
split.cpp:119:11: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  119 |     if (n == r.size()) return;
      |         ~~^~~~~~~~~~~
split.cpp: In function 'std::vector<int> find_split(int, int, int, int, std::vector<int>, std::vector<int>)':
split.cpp:138:23: warning: comparison of integer expressions of different signedness: 'int' and 'const size_t' {aka 'const long unsigned int'} [-Wsign-compare]
  138 |     for (int i = 0; i < M; i++) GRAPH::emplace_edge(p[i], q[i]);
      |                     ~~^~~
split.cpp:228:36: warning: comparison of integer expressions of different signedness: 'std::vector<int>::size_type' {aka 'long unsigned int'} and 'int' [-Wsign-compare]
  228 |   while(!que.empty() && avv.size() < targetSize[0].first) {
      |                         ~~~~~~~~~~~^~~~~~~~~~~~~~~~~~~~~

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