Submission #1014572

#	Time	Username	Problem	Language	Result	Execution time	Memory
1014572		NintsiChkhaidze	Mergers (JOI19_mergers)	C++17	48 / 100	3075 ms	49852 KiB

mergers

#include <bits/stdc++.h>
#define pb push_back
#define s second
#define f first
#define pb push_back
#define pii pair <int,int>
#define ll long long
#define int ll
using namespace std;

const int N = 3e5 + 5,inf = 1e18;
vector <pii> v[N],vec[N];
int col[N],tin[N],timer,d[25][N],dep[N];
int edge[N],val[N],n,D,I;

void dfs(int x,int par){
	d[0][x] = par;
	dep[x] = dep[par] + 1;
	tin[x] = ++timer;

	for (auto [to, id]: v[x]){
		if (to != par) {
			edge[to] = id;
			dfs(to,x);
		}
	}
}

int lca(int x,int y){
	if (dep[x] < dep[y ]) swap(x,y);
	for (int i = 20; i >= 0; i--)
		if (dep[d[i][x]] >= dep[y]) x = d[i][x];
	if (x == y) return x;

	for (int i = 20; i >= 0; i--)
		if (d[i][x] != d[i][y]) x = d[i][x],y = d[i][y];

	return d[0][x];
}

void getpath(int x,int y){
	int c = lca(x,y);
	while (x != c){
		val[edge[x]] = 0;
		x = d[0][x];
	}

	while (y != c){
		val[edge[y]] = 0;
		y = d[0][y];
	}
}	


void DFS(int x,int par,int dd){
	if (dd >= D) D = dd,I = x;

	for (auto [to,id]: v[x]){
		if (to != par) DFS(to,x,dd + val[id]);
	}
}

signed main() {
	ios_base::sync_with_stdio(0),cin.tie(NULL),cout.tie(NULL);

	int k;
	cin>>n>>k;

	for (int i = 1; i < n; i++){
		int a,b;
		cin>>a>>b;
		v[a].pb({b,i});
		v[b].pb({a,i});
		val[i] = 1;
	}

	dfs(1,1);
	for (int j = 1; j <= 20 ; j++)
		for (int i = 1; i <= n; i++)
			d[j][i] = d[j - 1][d[j - 1][i]];

	for (int i = 1; i <= n; i++)
		cin >> col[i],vec[col[i]].pb({tin[i],i});

	for (int i = 1; i <= k; i++){
		sort(vec[i].begin(),vec[i].end());
		for (int j = 1; j < vec[i].size(); j++){
			getpath(vec[i][j].s,vec[i][j - 1].s);
		}
	}

	int sum = 0;
	for (int i = 1; i < n; i++)
		sum += val[i];

	int ans = 0;
	while (sum > 0){
		++ans;
		I = D = 0; DFS(1,1,0);
		int x1 = I; 
		I = D = 0; DFS(x1,x1,0);
		int y1 = I;

		getpath(x1,y1);
		sum -= D;
	}

	cout<<ans;
}

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:87:21: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<std::pair<long long int, long long int> >::size_type' {aka 'long unsigned int'} [-Wsign-compare]
   87 |   for (int j = 1; j < vec[i].size(); j++){
      |                   ~~^~~~~~~~~~~~~~~

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