Submission #570301

#	Time	Username	Problem	Language	Result	Execution time	Memory
570301		franfill	Mergers (JOI19_mergers)	C++17	0 / 100	101 ms	9572 KiB

mergers

#include<bits/stdc++.h>
using namespace std;

struct dsu
{
	vector < int > pa;
	vector < int > siz;

	dsu(int n)
	{
		pa.resize(n);
		siz.resize(n, 1);
		iota(pa.begin(), pa.end(), 0);
	}
	
	dsu(int n, vector < int > &sizes) : siz(sizes)
	{
		pa.resize(n);
		iota(pa.begin(), pa.end(), 0);
	}

	int find(int x)
	{
		return pa[x]==x?x:(pa[x]=find(pa[x]));
	}

	bool unite(int a, int b)
	{
		a=find(a), b=find(b);
		if (a == b) return false;
		if (siz[a] < siz[b]) swap(a, b);
		siz[a] += siz[b];
		pa[b] = a;
		return true;
	}

	int get_siz(int x)
	{
		return siz[find(x)];
	}
};

int main()
{
	int n, k;
	cin >> n >> k;
	vector < vector < int > > g(n);
	for (int i = 0; i < n-1; i++)
	{
		int a, b;
		cin >> a >> b;
		a--, b--;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	vector < int > state(n), sizes(k, 0);
	for (auto &s : state)
	{
		cin >> s; s--;
		sizes[s]++;
	}
	dsu cities(n), states(k, sizes);	
	vector < int > conn(n, 0);
	int ans = 0;
	auto dfs = [&] (int x, int p, auto dfs) -> bool
	{
		for (auto y : g[x]) if (y != p)
		{
			if (dfs(y, x, dfs))
			{
				conn[x] += conn[y]-1;
				cities.unite(x, y);
				states.unite(state[x], state[y]);
			}
			else
			{
				conn[x]++;
			}
		}

		if (cities.get_siz(x) == states.get_siz(state[x]))
		{
			if (p != -1) conn[x]++;
			if (conn[x] == 1) ans++;
			return false;
		}
		else
		{
			return true;
		}
	};
	dfs(0, -1, dfs);
	cout << max(0, ans-1) << "\n";
}

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