Submission #376812

#	Time	Username	Problem	Language	Result	Execution time	Memory
376812		8e7	Mergers (JOI19_mergers)	C++14	100 / 100	1128 ms	119492 KiB

mergers

//Challenge: Accepted
#include <iostream>
#include <algorithm>
#include <vector>
#include <utility>
#define ll long long
#define maxn 500005
#define pii pair<int, int>
#define ff first
#define ss second
#define io ios_base::sync_with_stdio(0);cin.tie(0);
using namespace std;
vector<int> adj[maxn], tree[maxn];
int tot;
int col[2 * maxn], prv[maxn], cnt[maxn], par[maxn], ind[maxn];
//cnt: number of nodes colored i, ind: dsu parent
int num[maxn], ord[2 * maxn], arr[2 * maxn];
//ord: euler tour array, num: ans to query, arr:previous item
bool poss[maxn];

struct BIT {
	int bit[2 * maxn];
	void modify(int ind, int val) {
		if (ind == 0) return;
		for (;ind < 2 * maxn;ind += ind & (-ind)) bit[ind] += val;
	}
	int query(int ind) {
		int ret = 0;
		for (;ind > 0;ind -= ind & (-ind)) ret += bit[ind];
		return ret;
	}
} bb;

struct query{
	int l, r, id; //[l, r]
	query() {
		l = 0, r = 0, id = 0;
	}
	query(int a, int b, int c) {
		l = a, r = b, id = c;
	}
};
vector<query> que;
inline bool cmp(query a, query b) {
	return a.r < b.r;
}

int cur, ti = 1;
void dfs(int n, int pa) {
	ord[ti] = ord[ti + tot] = n;
	int lef = ti++;

	par[n] = pa;
	for (int v:adj[n]) {
		if (v != pa) {
			dfs(v, n);
		}
	}
	int rig = ti;
	//cout << n << " " << lef << " " << rig << endl;
	que.push_back(query(lef, rig - 1, n));
	que.push_back(query(rig, lef + tot - 1, n));

}

int find(int a) {
	return a == ind[a] ? a : ind[a] = find(ind[a]);
}
void Union(int a, int b) {
	ind[find(a)] = find(b);
}

int ans = 0;
void dfs2(int n, int par) {
	int chi = 0;
	for (int v:tree[n]) {
		if (v != par) {
			dfs2(v, n);
			chi++;
		}
	}
	if (chi == 0 || (n == 1 && chi == 1)) ans++;
}

int main() {
	io
	int n, k;
	cin >> n >> k;
	tot = n;
	for (int i = 0;i < n - 1;i++) {
		int a, b;
		cin >> a >> b;
		adj[a].push_back(b);
		adj[b].push_back(a);
	}
	for (int i = 1;i <= n;i++) {
		cin >> col[i];
		col[i + n] = col[i];
		cnt[col[i]]++;
		ind[i] = i;
	}


	dfs(1, 0);
	for (int i = 1;i <= 2 * n;i++) {
		arr[i] = prv[col[ord[i]]];
		//cout << arr[i] << " ";
		prv[col[ord[i]]] = i;
	}
	//cout << endl;
	sort(que.begin(), que.end(), cmp);
	int id = 0;
	for (int i = 1;i <= 2 * n;i++) {
		bb.modify(arr[i] + 1, 1);
		while (id < que.size() && que[id].r == i) {
			int val = que[id].r - que[id].l + 1 - (i - bb.query(que[id].l));
			//cout << que[id].l << " " << que[id].r << "  " << val << " " << (i - bb.query(que[id].l)) << endl;
			num[que[id].id] += val;
			id++;
		}
	}
	for (int i = 2;i <= n;i++) {
		//cout << num[i] << " ";
		if (num[i] > k) {
			poss[i] = true;
		}
	}
	//cout << endl;


	for (int i = 1;i <= n;i++) {
		if (poss[i]) Union(i, par[i]);
	}
	int done = 1;
	for (int i = 1;i <= n;i++) {
		if (par[i] != 0 && !poss[i]) {
			done = 0;
			int a = find(i), b = find(par[i]);
			tree[a].push_back(b);
			tree[b].push_back(a);
		}
	}
	if (done) cout << 0 << "\n";
	else {
		dfs2(find(1), 0);
		cout << (ans + 1) / 2 << "\n";
	}
}
/*
5 4
1 2
2 3
3 4
3 5
1
2
1
3
4

5 4
1 2
2 3
3 4
4 5
1
2
3
4
1

2 2
1 2
1
2

 */

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:115:13: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<query>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  115 |   while (id < que.size() && que[id].r == i) {
      |          ~~~^~~~~~~~~~~~

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