Submission #490532

#	Time	Username	Problem	Language	Result	Execution time	Memory
490532		Samboor	Capital City (JOI20_capital_city)	C++17	0 / 100	298 ms	90540 KiB

capital_city

// Grzegorz Ryn
// V LO Kraków

#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<int, int> pii;
typedef pair<ll, ll> pll;
typedef vector<bool> vb;
typedef vector<int> vi;
typedef vector<vector<int>> vvi;
typedef vector<ll> vl;
typedef vector<vector<ll>> vvl;

#define pb push_back
#define eb emplace_back
#define mp make_pair
#define mt make_tuple
#define st first
#define nd second
#define FOR(__VAR, __START, __END) for(int __VAR=__START; __VAR<=__END; __VAR++)
#define FORB(__VAR, __START, __END) for(int __VAR=__START; __VAR>=__END; __VAR--)
#define FORK(__VAR, __START, __END, __DIFF) for(int __VAR=__START; __VAR<=END; __VAR+=__DIFF)
#define all(__VAR) (__VAR).begin(), (__VAR).end()
#define rall(__VAR) (__VAR).rbegin(), (__VAR).rend()
#define DEBUG(__VAR) cout << #__VAR << ": " << __VAR << endl;

template<typename __T1, typename __T2>
ostream & operator<<(ostream &out, pair<__T1, __T2> &__VAR) {
	cout << "[" << __VAR.st << ", " << __VAR.nd << "]";
	return out;
}

template<typename __T>
ostream & operator<<(ostream &out, vector<__T> &__VAR) {
	cout << "[";
	FOR(i, 0, (int)__VAR.size()-2)
		cout << __VAR[i] << ", ";
	if(__VAR.size()>0)
		cout << __VAR[__VAR.size()-1];
	cout << "]" << endl;
	
	return out;
}

const int INF=1e9;

int N, K, n = 1<<18;
vector<vector<int>> g, inv_g, tree;
vector<int> pre_order, c, siz, heavy, depth, head, parent;
vector<bool> vis;

void prep_hld(int v) {
	vis[v] = true;
	for(int u:tree[v]) {
		if(vis[u]) {
			continue;
		}
		depth[u] = depth[v]+1;
		parent[u] = v;
		prep_hld(u);
		siz[v] += siz[u];
		if(heavy[v] == -1 || siz[heavy[v]] < siz[u]) {
			heavy[v] = u;
		}
	}
}

void get_ind(int v, int h) {
	static int next_it = 0;
	vis[v] = true;
	pre_order[v] = next_it++;
	head[v] = h;
	
	if(heavy[v] != -1) {
		get_ind(heavy[v], h);
	}
	for(int u:tree[v]) {
		if(vis[u]) {
			continue;
		}
		get_ind(u, u);
	}
}

void get_input() {
	cin >> N >> K;
	tree.resize(N);
	c.resize(N);
	pre_order.resize(N);
	parent.resize(N);
	for(int i=0; i<N-1; i++) {
		int a, b;
		cin >> a >> b;
		a--, b--;
		tree[a].pb(b);
		tree[b].pb(a);
	}
	for(int i=0; i<N; i++) {
		cin >> c[i];
		c[i]--;
	}
	
	vis.resize(N, false);
	siz.resize(N, 1);
	heavy.resize(N, -1);
	head.resize(N, 0);
	depth.resize(N, 0);
	
	prep_hld(0);
	fill(all(vis), false);
	get_ind(0, 0);
	
	g.resize(2*n);
	for(int i=n-1; i>0; i--) {
		g[i].pb(2*i);
		g[i].pb(2*i+1);
	}
}

void add(int l, int r, int v, int p, int q, int u) {
	if(r<p || q<l) {
		return;
	}
	if(p<=l && r<=q) {
		g[u].pb(v);
		return;
	}
	int mid = (l+r)/2;
	add(l, mid, 2*v, p, q, u);
	add(mid+1, r, 2*v+1, p, q, u);
}

int lca(int v, int u) {
	while(head[u] != head[v]) {
		if(depth[head[u]] < depth[head[v]]) {
			swap(u, v);
		}
		u = parent[head[u]];
	}
	if(depth[u] > depth[v]) {
		swap(u, v);
	}
	return u;
}

void update(int u, int v, int source) {
	while(head[u] != head[v]) {
		add(0, n-1, 1, pre_order[head[u]], pre_order[u], pre_order[source]);
		u = parent[head[u]];
	}
	add(0, n-1, 1, pre_order[u], pre_order[v], pre_order[source]);
}

stack<int> s;

void dfs1(int v) {
	vis[v] = true;
	for(int u:g[v]) {
		if(vis[u]) {
			continue;
		}
		dfs1(u);
	}
	s.push(v);
}

vector<int> comp;
int next_comp = 0;

void dfs2(int v) {
	vis[v] = true;
	comp[v] = next_comp;
	for(int u:inv_g[v]) {
		if(vis[u]) {
			continue;
		}
		dfs2(u);
	}
}

int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	
	get_input();
	
	vector<int> color_lca(K);
	for(int i=0; i<N; i++) {
		color_lca[c[i]] = i;
	}
	for(int i=0; i<N; i++) {
		color_lca[c[i]] = lca(color_lca[c[i]], i);
	}
	
	for(int i=0; i<N; i++) {
		update(i, color_lca[c[i]], pre_order[i]);
	}
	
	inv_g.resize(2*n);
	for(int i=0; i<2*n; i++) {
		for(int u:g[i]) {
			inv_g[u].pb(i);
		}
	}	
	
	vis.resize(2*n);
	fill(all(vis), false);
	for(int i=0; i<2*n; i++) {
		if(vis[i]) {
			continue;
		}
		dfs1(i);
	}
	
	fill(all(vis), false);
	comp.resize(2*n);
	while(!s.empty()) {
		int v=s.top();
		s.pop();
		if(vis[v]) {
			continue;
		}
		dfs2(v);
	}

	return 0;
}

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