Submission #376852

#TimeUsernameProblemLanguageResultExecution timeMemory
376852Kevin_Zhang_TWMergers (JOI19_mergers)C++17
100 / 100
2522 ms138528 KiB
#include <bits/stdc++.h>
using namespace std;
using ll = long long ;
#define pb emplace_back
#define AI(i) begin(i), end(i)
template<class T> bool chmin(T &a, T b) { return b < a && (a = b, true); }
template<class T> bool chmax(T &a, T b) { return a < b && (a = b, true); }
#ifdef KEV
#define DE(args...) kout("[ " + string(#args) + " ] = ", args)
void kout() { cerr << endl; }
template<class T, class ...U> void kout(T a, U ...b) { cerr << a << ' ', kout(b...); }
template<class T> void debug(T l, T r) { while (l != r) cerr << *l << " \n"[next(l) == r], ++l; }
#else
#define DE(...) 0
#define debug(...) 0
#endif
const int MAX_N = 500010, MAX_K = 20;

int n, k, belong[MAX_N];
vector<int> edge[MAX_N];
vector<int> state[MAX_N];
ll sum[MAX_N];

int anc[MAX_K][MAX_N], in[MAX_N], out[MAX_N];
bool isanc(int a, int b) { return in[a] <= in[b] && out[a] >= out[b]; }
void dfs(int x, int lst = 1) {
	static int t;
	anc[0][x] = lst;
	in[x] = ++t;
	for (int u : edge[x]) if (u != lst)
		dfs(u, x);
	out[x] = t;
}

void dfs2(int x, int lst = -1) {
	for (int u : edge[x]) if (u != lst) {
		dfs2(u, x);
		sum[x] += sum[u];
	}
}


vector<int> vir;
void dfs3(int x, int lst = -1) {

	int son = 0;
	for (int u : edge[x]) if (u != lst)
		if (sum[u] == 0) ++son;

	if (sum[x] == 0 && lst != -1) ++son;

	if (son) vir.pb(x);

	for (int u : edge[x]) if (u != lst)
		dfs3(u, x);
}
int getlca(int a, int b) {
	for (int d = MAX_K-1;d >= 0;--d)
		if (!isanc(anc[d][a], b))
			a = anc[d][a];
	return isanc(a, b) ? a : anc[0][a];
}

int32_t main() {
	ios_base::sync_with_stdio(0), cin.tie(0);
	cin >> n >> k;
	for (int a, b, i = 1;i < n;++i) {
		cin >> a >> b;
		edge[a].pb(b), edge[b].pb(a);
	}
	for (int i = 1;i <= n;++i) {
		cin >> belong[i];
		state[ belong[i] ].pb(i);
	}

	if (k == 1) return cout << 0 << '\n', 0;

	int rt = 1;

	dfs(rt);
	for (int d = 1;d < MAX_K;++d)
		for (int i = 1;i <= n;++i)
			anc[d][i] = anc[d-1][ anc[d-1][i] ];

	for (int i = 1;i <= k;++i) {
		if (state[i].size() > 1) {
			sort(AI(state[i]), [&](int a, int b) {
					return in[a] < in[b];
					});
			for (int x : state[i])
				++sum[x];
			sum[ getlca(state[i][0], state[i].back()) ] -= state[i].size();
		}
	}

	dfs2(rt);

	dfs3(rt);

	sort(AI(vir), [&](int a, int b) {
			return in[a] < in[b]; });

	vector<int> all;

	for (int i = 0;i < vir.size();++i) {
		all.pb(vir[i]);
		if (i+1 < vir.size())
			all.pb(getlca(vir[i], vir[i+1]));
   	}

	sort(AI(all), [&](int a, int b) {
			return in[a] < in[b]; });

	all.erase(unique(AI(all)), end(all));

	debug(AI(all));



	vector<int> stk;

	vector<int> cnt(n + 10);

	for (int u : all) {
		while (stk.size()) {
			int x = stk.back();
			DE(x, u);
			if (isanc(x, u)) break;
			stk.pop_back();
		}
		if (stk.size()) {
			int x = stk.back();
			++cnt[x], ++cnt[u];
		}
		stk.pb(u);
	}

	assert(all.size() != 1);

	int res = 0;

	for (int u : all)
		res += cnt[u] == 1;

	res = (res+1) / 2;

	cout << res << '\n';
}

Compilation message (stderr)

mergers.cpp: In function 'int32_t main()':
mergers.cpp:105:19: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  105 |  for (int i = 0;i < vir.size();++i) {
      |                 ~~^~~~~~~~~~~~
mergers.cpp:107:11: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  107 |   if (i+1 < vir.size())
      |       ~~~~^~~~~~~~~~~~
mergers.cpp:15:20: warning: statement has no effect [-Wunused-value]
   15 | #define debug(...) 0
      |                    ^
mergers.cpp:116:2: note: in expansion of macro 'debug'
  116 |  debug(AI(all));
      |  ^~~~~
mergers.cpp:14:17: warning: statement has no effect [-Wunused-value]
   14 | #define DE(...) 0
      |                 ^
mergers.cpp:127:4: note: in expansion of macro 'DE'
  127 |    DE(x, u);
      |    ^~
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...