Submission #110005

#	Time	Username	Problem	Language	Result	Execution time	Memory
110005		luciocf	Mergers (JOI19_mergers)	C++17	100 / 100	1585 ms	77532 KiB

mergers

#include <bits/stdc++.h>

using namespace std;

const int maxn = 5e5+10;

int n, k;

int pai[maxn], peso[maxn];

int s[maxn], tot[maxn];

int in[maxn], out[maxn], val[maxn], sz[maxn], tt;
int qtdZero, qtdTot, sub[maxn];

int grau[maxn];

vector<int> grafo[maxn];

void Init(void)
{
	for (int i = 1; i <= n; i++)
		pai[i] = i, peso[i] = 1;
}

int Find(int x)
{
	if (pai[x] == x) return x;
	return pai[x] = Find(pai[x]);
}

void Join(int x, int y)
{
	x = Find(x), y = Find(y);
	if (x == y) return;

	if (peso[x] < peso[y]) swap(x, y);

	pai[y] = x, peso[x] += peso[y];
}

void dfs(int u, int p)
{
	sz[u] = 1;
	in[u] = ++tt, val[tt] = s[u];

	for (auto v: grafo[u])
	{
		if (v == p) continue;

		dfs(v, u);

		sz[u] += sz[v];
	}

	out[u] = tt;
}

void solve(int u, int p, bool heavy)
{
	int maior = -1, ind = -1;

	for (auto v: grafo[u])
		if (v != p && sz[v] > maior)
			maior = sz[v], ind = v;

	for (auto v: grafo[u])
		if (v != p && v != ind)
			solve(v, u, 0);

	if (ind != -1)
		solve(ind, u, 1);

	--sub[s[u]];

	if (sub[s[u]] == 0)
		++qtdZero;

	if (sub[s[u]] == tot[s[u]]-1)
		--qtdTot;

	for (auto v: grafo[u])
	{
		if (v != p && v != ind)
		{
			for (int i = in[v]; i <= out[v]; i++)
			{
				--sub[val[i]];
				
				if (sub[val[i]] == 0)
					++qtdZero;

				if (sub[val[i]] == tot[val[i]]-1)
					--qtdTot;
			}
		}
	}

	if (qtdZero+qtdTot < k)
		Join(u, p);

	if (!heavy)
	{
		for (int i = in[u]; i <= out[u]; i++)
		{
			++sub[val[i]];
			
			if (sub[val[i]] == 1)
				--qtdZero;

			if (sub[val[i]] == tot[val[i]])
				++qtdTot;
		}
	}
}

int main(void)
{
	scanf("%d %d", &n, &k);

	for (int i = 1; i < n; i++)
	{
		int u, v;
		scanf("%d %d", &u, &v);

		grafo[u].push_back(v);
		grafo[v].push_back(u);
	}

	for (int i = 1; i <= n; i++)
	{
		scanf("%d", &s[i]);

		tot[s[i]]++;
		sub[s[i]]++;
	}

	qtdTot = k;

	Init();

	dfs(1, 0);
	solve(1, 0, 0);

	for (int i = 1; i <= n; i++)
		for (auto v: grafo[i])
			if (Find(v) != Find(i))
				grau[Find(v)]++;

	int leaf = 0;

	for (int i = 1; i <= n; i++)
		if (grau[i] == 1)
			leaf++;

	printf("%d\n", (leaf+1)/2);
}

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:119:7: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
  scanf("%d %d", &n, &k);
  ~~~~~^~~~~~~~~~~~~~~~~
mergers.cpp:124:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d %d", &u, &v);
   ~~~~~^~~~~~~~~~~~~~~~~
mergers.cpp:132:8: warning: ignoring return value of 'int scanf(const char*, ...)', declared with attribute warn_unused_result [-Wunused-result]
   scanf("%d", &s[i]);
   ~~~~~^~~~~~~~~~~~~

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