Submission #1262564

#	Time	Username	Problem	Language	Result	Execution time	Memory
1262564		PlayVoltz	Capital City (JOI20_capital_city)	C++20	1 / 100	96 ms	40376 KiB

capital_city

#include <cstdio>
#include <stdio.h>
#include <stdbool.h>
#include <iostream>
#include <map>
#include <vector>
#include <climits>
#include <stack>
#include <string>
#include <queue>
#include <algorithm>
#include <set>
#include <unordered_set>
#include <unordered_map>
#include <cmath>
#include <cctype>
#include <bitset>
#include <iomanip>
#include <cstring>
#include <numeric>
#include <cassert>
#include <random>
#include <chrono>
#include <fstream>
using namespace std;

#define int long long
#define pii pair<int, int>
#define mp make_pair
#define pb push_back
#define fi first
#define se second

int ans=LLONG_MAX/2;
vector<bool> del, visited;
vector<int> sz, got, cc, vect, full, par;
vector<vector<int> > graph, col;

void dfs_size(int node, int p){
	sz[node]=1;
	for (auto num:graph[node])if (num!=p&&!del[num])dfs_size(num, node), sz[node]+=sz[num];
}

int findcentroid(int node, int p, int size){
	for (auto num:graph[node])if (num!=p&&!del[num]&&sz[num]>size/2)return findcentroid(num, node, size);
	return node;
}

void dfs2(int node, int p, int c){
	par[node]=p;
	if (vect[node]==c)got.pb(node);
	full.pb(node);
	++cc[vect[node]];
	for (auto num:graph[node])if (num!=p&&!del[num])dfs2(num, node, c);
}

void dfs(int node){
	dfs_size(node, -1);
	node=findcentroid(node, -1, sz[node]);
	dfs2(node, -1, vect[node]);
	del[node]=1;
	if (col[vect[node]].size()==cc[vect[node]]){
		int res=0;
		visited[vect[node]]=1;
		while (got.size()){
			int c=par[got.back()];
			got.pop_back();
			if(c!=-1&&!visited[vect[c]]){
				if (col[vect[c]].size()!=cc[vect[c]])goto end;
				visited[vect[c]]=1;
				for (auto a:col[vect[c]])got.pb(a);
				++res;
			}
		}
		ans=min(ans, res);
	}
	end:;
	for (auto a:full)cc[vect[a]]=0, visited[a]=0;
	full.clear();
	got.clear();
}

int32_t main(){
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	int n, k, a, b;
	cin>>n>>k;
	par.resize(n+1);
	cc.resize(n+1, 0);
	sz.resize(n+1);
	del.resize(n+1, 0);
	col.resize(k+1);
	vect.resize(n+1);
	graph.resize(n+1);
	visited.resize(k+1, 0);
	for (int i=1; i<n; ++i){
		cin>>a>>b;
		graph[a].pb(b);
		graph[b].pb(a);
	}
	for (int i=1; i<=n; ++i)cin>>vect[i], col[vect[i]].pb(i);
	dfs(1);
	cout<<ans;
}

• Subtask 1: 1 / 1
• Subtask 2: 0 / 10
• Subtask 3: 0 / 30
• Subtask 4: 0 / 59