Submission #105567

#	Time	Username	Problem	Language	Result	Execution time	Memory
105567		Pro_ktmr	Mergers (JOI19_mergers)	C++14	100 / 100	1993 ms	102648 KiB

mergers

#include"bits/stdc++.h"
using namespace std;
#define LL long long
#define MP make_pair
#define PB push_back

struct UnionFind{
private:
	int N;
	vector<int> par;
public:
	void init(int n){
		N = n;
		par.clear();
		for(int i=0; i<N; i++) par.PB(i);
	}
	int root(int i){
		if(par[i] == i) return i;
		return par[i] = root(par[i]);
	}
	void unite(int i, int j){
		i = root(i);
		j = root(j);
		if(i == j) return;
		par[j] = i;
	}
	bool same(int i, int j){
		return (root(i) == root(j));
	}
};

int N,K;
pair<int,int> E[500000];
vector<int> edge[500000];
vector<int> node[500000];

int par[500000];
bool visit[500000] = {};
void dfs(int now){
	visit[now] = true;
	for(int i=0; i<edge[now].size(); i++){
		if(visit[edge[now][i]]) continue;
		par[edge[now][i]] = now;
		dfs(edge[now][i]);
	}
}
int d[500000] = {};
int depth(int i){
	if(i == 0) return 0;
	if(d[i] != 0) return d[i];
	return d[i] = depth(par[i])+1;
}

UnionFind uf;
vector<int> newEdge[500000];
int main(){
	cin >> N >> K;
	for(int i=0; i<N-1; i++){
		int A,B;
		cin >> A >> B;
		A--; B--;
		edge[A].PB(B);
		edge[B].PB(A);
		E[i] = MP(A,B);
	}
	par[0] = 0;
	dfs(0);
	for(int i=0; i<N; i++){
		int S;
		cin >> S;
		S--;
		node[S].PB(i);
	}
	
	//
	
	uf.init(N);
	for(int i=0; i<K; i++){
		if(node[i].size() == 0) continue;
		int now = uf.root(node[i][0]);
		for(int j=1; j<node[i].size(); j++){
			int tmp = uf.root(node[i][j]);
			while(now != tmp){
				if(depth(now) > depth(tmp)){
					uf.unite(par[now],now);
					now = uf.root(now);
				}
				else{
					uf.unite(par[tmp],tmp);
					tmp = uf.root(tmp);
				}
			}
		}
	}
	
	//
	
	for(int i=0; i<N-1; i++){
		if(uf.same(E[i].first, E[i].second)) continue;
		newEdge[uf.root(E[i].first)].PB(uf.root(E[i].second));
		newEdge[uf.root(E[i].second)].PB(uf.root(E[i].first));
	}
	int ans = 0;
	for(int i=0; i<N; i++){
		if(newEdge[i].size() == 1) ans++;
	}
	cout << (ans+1)/2 << endl;
	
	return 0;
}

Compilation message (stderr)

mergers.cpp: In function 'void dfs(int)':
mergers.cpp:41:16: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
  for(int i=0; i<edge[now].size(); i++){
               ~^~~~~~~~~~~~~~~~~
mergers.cpp: In function 'int main()':
mergers.cpp:81:17: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
   for(int j=1; j<node[i].size(); j++){
                ~^~~~~~~~~~~~~~~

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