Submission #263384

#	Time	Username	Problem	Language	Result	Execution time	Memory
263384		mjkocijan	Mergers (JOI19_mergers)	C++14	100 / 100	1923 ms	194492 KiB

mergers

#include <bits/stdc++.h>
using namespace std;
#define X first
#define Y second
#define pb push_back
typedef long long ll;
typedef pair<ll, ll> ii;
#define MAXN 500500
#define LOG 20

int n, k;
vector<int> g[MAXN];
int s[MAXN], rt[MAXN], sz[MAXN], dep[MAXN], st1[MAXN], st2[MAXN];
vector<int> sr[MAXN];
int anc[LOG][MAXN];

int dfs(int cv, int rod, int dub = 0)
{
	dep[cv] = dub;
	anc[0][cv] = rod;
	st1[cv] = 1;
	for (int i: g[cv]) {
		if (rod != i) {
			st1[cv] += dfs(i, cv, dub + 1);
		}
	}
	return st1[cv];
}

int lca(int x, int y)
{
	if (dep[x] > dep[y]) swap(x, y);
	
	for (int i = LOG - 1; i >= 0; i--) {
		int novi = anc[i][y];
		if (dep[novi] >= dep[x])
			y = novi;
	}
	
	if (x == y) return x;
	
	for (int i = LOG - 1; i >= 0; i--) {
		if (anc[i][x] != anc[i][y]) {
			x = anc[i][x];
			y = anc[i][y];
		}
	}
	
	return anc[0][x];
}

int fn(int x)
{
	if (x == rt[x]) return x;
	
	return rt[x] = fn(rt[x]);
}

void merg(int x, int y)
{
	x = fn(x);
	y = fn(y);
	
	if (x == y) return;
	
	if (dep[x] < dep[y]) return;

	//if (sz[x] > sz[y]) swap(x, y);
	
	rt[x] = y;
	//sz[y] += sz[x] == sz[y];
	st2[y] += st2[x];
}

//set<int> s;

set<int> sus[MAXN];

int main()
{
	scanf("%d%d", &n, &k);
	for (int i = 0; i < n - 1; i++) {
		int q, w;
		scanf("%d%d", &q, &w);
		q--; w--;
		g[q].pb(w);
		g[w].pb(q);
	}
	for (int i = 0; i < n; i++) {
		rt[i] = i;
		st2[i] = 1;
		scanf("%d", &s[i]);
		s[i]--;
		sr[s[i]].pb(i);
	}
	dfs(0, 0);
	for (int i = 1; i < LOG; i++) {
		for (int j = 0; j < n; j++) {
			anc[i][j] = anc[i - 1][anc[i - 1][j]];
		}
	}
	for (int i = 0; i < k; i++) {
		int lcaa = sr[i][0];
		for (int j = 1; j < sr[i].size(); j++) {
			lcaa = lca(lcaa, sr[i][j]);
		}
		for (int j: sr[i]) {
			j = fn(j);
			while (dep[j] > dep[lcaa]) {
				merg(j, anc[0][j]);
				j = fn(j);
			}
		}
	}
	int reza = 0;
	for (int i = 0; i < n; i++) {
		for (int j: g[i]) {
			if (fn(i) != fn(j))
				sus[fn(i)].insert(fn(j));
		}
		//cout << i+1<< ' '<<fn(i)+1<<endl;
		//cout << i+1<<' '<<fn(i)+1<<' '<<st1[i]<<' '<<st2[i]<<endl;
		//if (i == fn(i) && st1[i] == st2[i]) reza++;
	}
	for (int i = 0; i < n; i++) {
		if (sus[i].size() == 1)
			reza++;
		//cout << i+1<<' '<<
	}
	/*if (reza == 1)
		printf("0\n");
	else*/
		printf("%d\n", (reza + 1) / 2);
  return 0;
}

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:104:21: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  104 |   for (int j = 1; j < sr[i].size(); j++) {
      |                   ~~^~~~~~~~~~~~~~
mergers.cpp:81:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   81 |  scanf("%d%d", &n, &k);
      |  ~~~~~^~~~~~~~~~~~~~~~
mergers.cpp:84:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   84 |   scanf("%d%d", &q, &w);
      |   ~~~~~^~~~~~~~~~~~~~~~
mergers.cpp:92:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   92 |   scanf("%d", &s[i]);
      |   ~~~~~^~~~~~~~~~~~~

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