Submission #1039907

#	Time	Username	Problem	Language	Result	Execution time	Memory
1039907		ymm	Capital City (JOI20_capital_city)	C++17	11 / 100	1216 ms	524288 KiB

capital_city

#include <bits/stdc++.h>
#define Loop(x,l,r) for (ll x = (l); x < (ll)(r); ++x)
#define LoopR(x,l,r) for (ll x = (r)-1; x >= (ll)(l); --x)
typedef long long ll;
typedef std::pair<int, int> pii;
typedef std::pair<ll , ll > pll;
using namespace std;

const int N = 200'010;
const int lg = 20;

struct snode;
struct gnode;

struct gnode {
	gnode *par;
	gnode *anc[lg];
	vector<gnode *> adj;
	gnode *hchild, *hroot, *hleaf;
	snode *sroot;
	int st, ft, ht;
	int col;

	int sz() { return ft - st; }
	int hsz() { return hleaf->ht - hroot->ht + 1; }
} graph[N];

struct snode {
	vector<snode *> adj;
	vector<snode *> adjt;
	int weight;
	int vis;

	static void add_edge(snode *v, snode *u) {
		v->adj.push_back(u);
		u->adjt.push_back(v);
	}
} pool[(1<<28) / sizeof(gnode)];

vector<snode *> all_snodes;

snode *new_snode()
{
	static int p = 0;
	all_snodes.push_back(&pool[p]);
	return &pool[p++];
}

void dfsg1(gnode *v, gnode *p, int &t, int h)
{
	v->par = p;
	v->st = t++;
	v->ht = h;

	v->anc[0] = v->par? v->par: v;
	Loop (i,0,lg-1)
		v->anc[i + 1] = v->anc[i]->anc[i];

	for (auto u : v->adj)
		if (u != p)
			dfsg1(u, v, t, h + 1);
	v->ft = t++;
	for (auto u : v->adj) {
		if (u == p)
			continue;
		if (v->hchild == 0 || v->hchild->sz() < u->sz())
			v->hchild = u;
	}
}

gnode *lca(gnode *v, gnode *u)
{
	if (v->ht < u->ht)
		swap(v, u);
	int diff = v->ht - u->ht;
	Loop (i,0,lg)
		if ((diff >> i) & 1)
			v = v->anc[i];
	if (v == u)
		return v;
	LoopR (i,0,lg) {
		if (v->anc[i] != u->anc[i]) {
			v = v->anc[i];
			u = u->anc[i];
		}
	}
	return v->anc[0];
}

gnode gn[N];
snode col_snode[N];
int n, k;

void init_col_snode()
{
	Loop (i,0,k) {
		col_snode[i].weight = 1;
		all_snodes.push_back(&col_snode[i]);
	}
}

snode *make_seg(snode **ptr, int l, int r)
{
	if (r - l == 1)
		return ptr[l];
	snode *sn = new_snode();
	snode::add_edge(sn, make_seg(ptr, l, (l + r)/2));
	snode::add_edge(sn, make_seg(ptr, (l + r)/2, r));
	return sn;
}

snode *make_seg(gnode *u, gnode *d)
{
	vector<snode *> vec(d->ht - u->ht + 1);
	gnode *v = d;
	Loop (i,0,vec.size()) {
		vec.end()[-i-1] = &col_snode[v->col];
		v = v->par;
	}
	return make_seg(vec.data() - u->ht, u->ht, d->ht + 1);
}

template<typename Fn>
void seg_apply(snode *v, Fn fn, int l, int r, int b, int e)
{
	// if (l <= b && e <= r) {
	// 	fn(v);
	// 	return;
	// }
	if (r <= b || e <= l)
		return;
	if (e - b == 1) {
		fn(v);
		return;
	}
	seg_apply(v->adj[0], fn, l, r, b, (b+e)/2);
	seg_apply(v->adj[1], fn, l, r, (b+e)/2, e);
}

void dfsg2(gnode *v, gnode *p, gnode *hroot)
{
	for (gnode *u : v->adj) {
		if (u == p)
			continue;
		dfsg2(u, v, u == v->hchild? hroot: u);
	}
	v->hleaf = v->hchild? v->hchild->hleaf: v;
	v->hroot = hroot;
	if (v == hroot)
		v->sroot = make_seg(v, v->hleaf);
}

template<typename Fn>
void path_apply(gnode *u, gnode *d, Fn fn)
{
	while (d->hroot->ht > u->ht) {
		gnode *h = d->hroot;
		seg_apply(h->sroot, fn, h->ht, d->ht + 1, h->ht, h->hleaf->ht + 1);
		d = h->par;
	}
	gnode *h = d->hroot;
	assert(h == u->hroot);
	seg_apply(h->sroot, fn, u->ht, d->ht + 1, h->ht, h->hleaf->ht + 1);
}

void path_add_col_to(gnode *u, gnode *d, snode *col_node)
{
	//cerr << col_node - col_snode << " => " << u - graph << " --- " << d - graph << '\n';
	path_apply(u, d, [=](snode *sn) {
		//cerr << col_node - col_snode << " -> " << sn - col_snode << '\n';
		snode::add_edge(col_node, sn);
	});
	// while (d && d->ht >= u->ht) {
	// 	snode::add_edge(col_node, &col_snode[d->col]);
	// 	//cerr << col_node - col_snode << " -> " << d->col << '\n';
	// 	d = d->par;
	// }
}

void scc_dfs1(snode *v, vector<snode *> &vec)
{
	v->vis = 1;
	for (snode *u : v->adj) {
		if (!u->vis)
			scc_dfs1(u, vec);
	}
	vec.push_back(v);
}

int scc_dfs2(snode *v, int viscol)
{
	v->vis = viscol;
	int ans = v->weight;
	for (snode *u : v->adjt) {
		if (u->vis == -1)
			ans += scc_dfs2(u, viscol);
	}
	return ans;
}

int scc_solve()
{
	vector<snode *> vec;
	for (snode *sn : all_snodes)
		sn->vis = 0;
	for (snode *sn : all_snodes) {
		if (!sn->vis)
			scc_dfs1(sn, vec);
	}
	reverse(vec.begin(), vec.end());
	for (snode *sn : all_snodes)
		sn->vis = -1;
	vector<int> scc_weight;
	for (snode *sn : vec) {
		if (sn->vis != -1)
			continue;
		int x = scc_dfs2(sn, scc_weight.size());
		scc_weight.push_back(x);
	}
	vector<bool> scc_ok(scc_weight.size(), true);
	//Loop (i,0,k)
	//	cerr << i << ": " << col_snode[i].vis << '\n';
	for (snode *v : all_snodes) for (snode *u : v->adj) {
		if (v->vis != u->vis)
			scc_ok[v->vis] = false;
	}
	int ans = INT_MAX;
	Loop (i,0,scc_weight.size()) {
		if (!scc_ok[i])
			continue;
		assert(scc_weight[i]);
		ans = min(ans, scc_weight[i]);
	}
	return ans;
}

vector<gnode *> gnodes_of_col[N];

void make_virtuals()
{
	auto cmp = [](gnode *v, gnode *u) { return v->st < u->st; };
	Loop (col,0,k) {
		vector<gnode *> vec = gnodes_of_col[col];
		sort(vec.begin(), vec.end(), cmp);
		if (vec.empty())
			continue;

		int cnt = vec.size();
		Loop (i,0,cnt-1)
			vec.push_back(lca(vec[i], vec[i + 1]));
		sort(vec.begin(), vec.end(), cmp);

		vector<gnode *> st;
		for (gnode *v : vec) {
			while (st.size() && st.back()->ft <= v->st)
				st.pop_back();
			if (st.size())
				path_add_col_to(st.back(), v, &col_snode[col]);
			st.push_back(v);
		}
	}
}

int main()
{
	cin.tie(0) -> sync_with_stdio(false);
	cin >> n >> k;
	Loop (i,0,n-1) {
		int v, u;
		cin >> v >> u;
		--v; --u;
		graph[v].adj.push_back(&graph[u]);
		graph[u].adj.push_back(&graph[v]);
	}
	Loop (i,0,n) {
		gnode *v = &graph[i];
		cin >> v->col;
		v->col--;
		gnodes_of_col[v->col].push_back(v);
	}
	init_col_snode();
	int dummy = 0;
	dfsg1(&graph[0], nullptr, dummy, 0);
	dfsg2(&graph[0], nullptr, &graph[0]);
	make_virtuals();
	int ans = scc_solve();
	cout << ans-1 << '\n';
}

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 0 / 30
• Subtask 4: 0 / 59