Submission #138031

#	Time	Username	Problem	Language	Result	Execution time	Memory
138031		rajarshi_basu	Mergers (JOI19_mergers)	C++14	0 / 100	110 ms	42404 KiB

mergers

#include <bits/stdc++.h>

#define FOR(i,n) for(int i=0;i<n;i++)
#define FORE(i,a,b) for(int i = a;i<=b;i++)
#define ll long long int
#define pb push_back
#define vi vector<int>
#define ff first
#define ss second
#define vv vector
#define ii pair<int,int>

using namespace std;

const int MAXN = 5e5 + 1;
const int LOGN = 20;

int n,k;


vi g[MAXN];
int color[MAXN];
vi verticesOfColor[MAXN];
int sparseTable[LOGN][MAXN];
int depth[MAXN];
int initialLCA[MAXN];
int degree[MAXN];

void dfs_setup(int node,int p = -1){
	sparseTable[0][node] = p;
	if(p == -1)depth[node] = -1;else depth[node] = depth[p] + 1;
	for(auto e : g[node]){
		if(e != p)dfs_setup(e,node);
	}
}

void computeSparseTable(){
	FOR(i,LOGN){
		if(i == 0)continue;
		FOR(node,n){
			int p = sparseTable[i-1][node];
			if(p == -1)sparseTable[i][node] = -1;
			else sparseTable[i][node] = sparseTable[i-1][p];
		}
	}
}

int LCA(int a,int b){
	if(depth[a] > depth[b])swap(a,b);
	for(int x = LOGN-1;x>=0;x--){
		if(sparseTable[x][b] != -1 and depth[sparseTable[x][b]] >= depth[a])
			b = sparseTable[x][b];
	}
	if(a == b)return a;
	for(int x = LOGN-1;x>=0;x--){
		if(sparseTable[x][a] != sparseTable[x][b]){
			a = sparseTable[x][a];
			b = sparseTable[x][b];
		}
	}
	return sparseTable[0][a];
}


struct DSU{
	int parent[MAXN];
	int ID[MAXN];
	void init(){
		FOR(i,n) parent[i] = i;
		FOR(i,n) ID[i] = i;
	}
	int find(int a){
		if(parent[a] != a)parent[a] = find(parent[a]);
		return parent[a];
	}
	void merge(int a,int b){
		int pa = find(a);
		int pb = find(b);
		parent[pa] = pb;
		if(depth[pa] < depth[pb])ID[pb] = pa;else ID[pb] = pa;
	}
	int id(int a){
		return ID[find(a)];
	}
	bool isSame(int a,int b){
		return find(a) == find(b);
	}
} dsu;

bool done[MAXN];

int main(){
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	cin >> n >> k;
	vv<ii> allE;
	FOR(i,n-1){
		int a,b;cin >> a >> b;a--;b--;
		g[a].pb(b);g[b].pb(a);
		allE.pb({a,b});
	}
	FOR(i,k)initialLCA[i] = -1;
	dfs_setup(0);
	computeSparseTable();
	FOR(i,n){
		int a;cin >> a;a--;
		color[i] = a;
		verticesOfColor[a].pb(i);
		//if(initialLCA[a] == -1)initialLCA[a] = i;
		//else initialLCA[a] = LCA(initialLCA[a],i);
	}
	FOR(i,k){
		if(verticesOfColor[i].size() == 1)continue;
		FORE(j,1,(int)verticesOfColor[i].size()-1){
			int node1 = dsu.find(verticesOfColor[i][j]);
			int node2 = dsu.find(verticesOfColor[i][j-1]);
			while(node1 != node2){
				if(depth[node1] > depth[node2])swap(node1,node2);
				int p = dsu.id(sparseTable[0][node2]);
				dsu.merge(p,node2);
				node2 = p;
			}
		}
	}
	
	for(auto e : allE)if(!dsu.isSame(e.ff,e.ss))degree[e.ff]++,degree[e.ss]++;
	int cnt = 0;
	FOR(i,n)if(degree[i] == 1)cnt++;
	cout << cnt << endl;
	return 0;
}

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