Submission #1014663

#	Time	Username	Problem	Language	Result	Execution time	Memory
1014663		NintsiChkhaidze	Mergers (JOI19_mergers)	C++17	70 / 100	2514 ms	262144 KiB

mergers

#include <bits/stdc++.h>
#define pb push_back
#define s second
#define f first
#define pb push_back
#define pii pair <int,int>
#define ll long long
#define int ll
using namespace std;

const int N = 5e5 + 5,inf = 1e18;
vector <pii> v[N],vec[N];
int col[N],tin[N],timer,d[25][N],dep[N];
int edge[N],val[N],n,D,I,p[N];
bool vis[N],del[N];
set <int> st[N];

void dfs(int x,int par){
	d[0][x] = par;
	dep[x] = dep[par] + 1;
	tin[x] = ++timer;

	for (auto [to, id]: v[x]){
		if (to != par) {
			edge[to] = id;
			dfs(to,x);
		}
	}
}

int lca(int x,int y){
	if (dep[x] < dep[y ]) swap(x,y);
	for (int i = 20; i >= 0; i--)
		if (dep[d[i][x]] >= dep[y]) x = d[i][x];
	if (x == y) return x;

	for (int i = 20; i >= 0; i--)
		if (d[i][x] != d[i][y]) x = d[i][x],y = d[i][y];

	return d[0][x];
}

void getpath(int x,int y){
	int c = lca(x,y);
	while (x != c){
		val[edge[x]] = 0;
		x = d[0][x];
	}

	while (y != c){
		val[edge[y]] = 0;
		y = d[0][y];
	}
}	


void DFS(int x,int par,int dd){
	if (dd >= D) D = dd,I = x;

	for (auto [to,id]: v[x]){
		if (to != par) DFS(to,x,dd + val[id]);
	}
}

int P(int x){
	if(x==p[x])return x;
	return p[x]=P(p[x]);
}

void dsu(int x,int y){
	int px = P(x),py = P(y);
	if (st[px].size() < st[py].size()) swap(px,py);
	p[py] = px;
	del[py] = 1;
	
	for (int v: st[py])
		st[px].insert(v);
	
	st[py].clear();
}

signed main() {
	ios_base::sync_with_stdio(0),cin.tie(NULL),cout.tie(NULL);

	int k;
	cin>>n>>k;

	for (int i = 1; i < n; i++){
		int a,b;
		cin>>a>>b;
		st[a].insert(b);
		st[b].insert(a);
		v[a].pb({b,i});
		v[b].pb({a,i});
		val[i] = 1;
	}

	dfs(1,1);
	for (int j = 1; j <= 20 ; j++)
		for (int i = 1; i <= n; i++)
			d[j][i] = d[j - 1][d[j - 1][i]];

	for (int i = 1; i <= n; i++)
		cin >> col[i],vec[col[i]].pb({tin[i],i});

	for (int i = 1; i <= k; i++){
		sort(vec[i].begin(),vec[i].end());
		for (int j = 1; j < vec[i].size(); j++){
			getpath(vec[i][j].s,vec[i][j - 1].s);
		}
	}

	for (int i=1;i<=n;i++)	
		p[i] = i;

	for (int i = 2; i <= n; i++)
		if (!val[edge[i]]) dsu(i,d[0][i]);

	for (int i = 1; i <= n;i++){
		if (del[i]) continue;
		set <int> nw;
		for (int v: st[i])
			nw.insert(P(v));

		if (nw.find(i) != nw.end()) nw.erase(nw.find(i));
		st[i] = nw;
	}

	int ans=0;
	for (int i = 1; i <= n; i++){
		if (del[i]) continue;
		ans += ((int)st[i].size() == 1);
	}
	cout<<(ans + 1)/2;
}

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:108:21: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  108 |   for (int j = 1; j < vec[i].size(); j++){
      |                   ~~^~~~~~~~~~~~~~~

• Subtask 1: 10 / 10
• Subtask 2: 24 / 24
• Subtask 3: 14 / 14
• Subtask 4: 22 / 22
• Subtask 5: 0 / 30