Submission #538152

#	Time	Username	Problem	Language	Result	Execution time	Memory
538152		Haruto810198	Capital City (JOI20_capital_city)	C++17	100 / 100	1640 ms	181876 KiB

capital_city

#include <bits/stdc++.h>
 
using namespace std;
 
#define int long long
#define double long double
 
#define FOR(i, l, r, d) for(int i=(l); i<=(r); i+=(d))
#define szof(x) ((int)(x).size())
 
#define vi vector<int>
#define pii pair<int, int>
 
#define F first
#define S second
 
#define pb push_back
#define eb emplace_back
#define mkp make_pair

const int INF = INT_MAX;
const int LNF = INF*INF;
const int MOD = 1000000007;
const double eps = 1e-12;
 
//#pragma GCC optimize("Ofast")
//#pragma GCC optimize("unroll-loops")

const int MAX = 200010;

int n, k;
vi edge_t[MAX];
int city[MAX];
vi occ[MAX];
int res;

// HLD
int par[MAX], sz[MAX], dep[MAX], heavy[MAX], dfn[MAX], fr[MAX];
int ts;

// segment tree optimization
#define V st[cidx]

struct Node{
	int lc, rc;
	int sl, sr;
};
int nodes = 1;
Node st[2*MAX];
int leaf[MAX]; // leaf[i] = id. of node [i, i] in st
int w[2*MAX];

// SCC
vi edge_g[2*MAX];
vi edge_g_rev[2*MAX];
vi dfstk;
vi cur_scc;
bool vis[2*MAX];
int scccnt;
int sccid[2*MAX];
vi edge_scc[2*MAX];
int scc_sz[2*MAX];

// ===============================================================

// HLD
void find_heavy(int v, int pv){
	par[v] = pv;
	sz[v] = 1;
	if(v == pv) dep[v] = 0;
	else dep[v] = dep[pv] + 1;
	
	int Max = 0;
	heavy[v] = -1;
	for(int i : edge_t[v]){
		if(i == pv) continue;
		find_heavy(i, v);
		sz[v] += sz[i];
		if(Max < sz[i]){
			Max = sz[i];
			heavy[v] = i;
		}
	}
}

void HLD(int v, int pv, int frv){
	dfn[v] = ++ts;
	fr[v] = frv;
	
	if(heavy[v] != -1) HLD(heavy[v], v, frv);

	for(int i : edge_t[v]){
		if(i == pv or i == heavy[v]) continue;
		HLD(i, v, i);
	}
}

vector<pii> Segs(int u, int v){
	vector<pii> ret;
	while(fr[u] != fr[v]){
		if(dep[fr[u]] > dep[fr[v]]) swap(u, v);
		ret.eb(dfn[fr[v]], dfn[v]);
		v = par[fr[v]];
	}

	if(dep[u] > dep[v]) swap(u, v);
	ret.eb(dfn[u], dfn[v]);

	return ret;
}

// segment tree
void build(int cidx, int cl, int cr){
	V.sl = cl;
	V.sr = cr;
	if(cl < cr){
		int mid = (cl + cr) / 2;
		
		V.lc = ++nodes;
		edge_g[cidx].pb(V.lc);
		build(nodes, cl, mid);

		V.rc = ++nodes;
		edge_g[cidx].pb(V.rc);
		build(nodes, mid+1, cr);
	}
	else{
		leaf[cl] = cidx;
	}
}

void add_edge(int cidx, int ml, int mr, int from){
	//if(cidx == 1) cerr<<"add_edge "<<from<<" -> ["<<ml<<", "<<mr<<"] "<<endl;
	// ml, mr, from : dfn of nodes
	if(mr < V.sl or V.sr < ml) return;
	if(ml <= V.sl and V.sr <= mr){
		edge_g[leaf[from]].pb(cidx);
		return;
	}

	add_edge(V.lc, ml, mr, from);
	add_edge(V.rc, ml, mr, from);
}

// SCC
void dfs1(int v){
	vis[v] = 1;
	for(int i : edge_g[v]){
		if(vis[i] == 0) dfs1(i);
	}
	dfstk.pb(v);
}

void dfs2(int v){
	vis[v] = 1;
	for(int i : edge_g_rev[v]){
		if(vis[i] == 0) dfs2(i);
	}
	cur_scc.pb(v);
}

void find_sccs(){

	FOR(v, 1, nodes, 1){
		if(vis[v] == 0) dfs1(v);
	}
	
	FOR(v, 1, nodes, 1){
		for(int i : edge_g[v]){
			edge_g_rev[i].pb(v);
		}
		vis[v] = 0;
	}
	
	while(!dfstk.empty()){
		int v = dfstk.back();
		dfstk.pop_back();

		if(vis[v] == 0){
			dfs2(v);
			scccnt++;
			for(int i : cur_scc){
				sccid[i] = scccnt;
				scc_sz[scccnt] += w[i];
			}
			cur_scc.clear();
		}
	}
	/*	
	cerr<<"sccid : ";
	FOR(v, 1, nodes, 1){
		cerr<<sccid[v]<<' ';
	}
	cerr<<endl;
	*/
	FOR(v, 1, nodes, 1){
		for(int i : edge_g[v]){
			if(sccid[v] != sccid[i]) edge_scc[sccid[v]].pb(sccid[i]);
		}
	}
	
	FOR(v, 1, scccnt, 1){
		auto it = unique(edge_scc[v].begin(), edge_scc[v].end());
		edge_scc[v].resize(distance(edge_scc[v].begin(), it));
		/*	
		cerr<<"scc_sz["<<v<<"] = "<<scc_sz[v]<<", ";
		cerr<<"edge_scc["<<v<<"] : ";
		for(int i : edge_scc[v]){
			cerr<<i<<' ';
		}
		cerr<<endl;
		*/
	}
	
	res = INF;
	FOR(v, 1, scccnt, 1){
		if(edge_scc[v].empty()) res = min(res, scc_sz[v]);
	}
}

signed main(){
	
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    cout.tie(0);
	
	cin>>n>>k;
	FOR(i, 1, n-1, 1){
		int u, v;
		cin>>u>>v;
		edge_t[u].pb(v);
		edge_t[v].pb(u);
	}

	FOR(i, 1, n, 1){
		cin>>city[i];
		occ[city[i]].pb(i);
	}

	// HLD
	find_heavy(1, 1);
	HLD(1, 1, 1);
	/*	
	cerr<<"dfn : ";
	FOR(i, 1, n, 1){
		cerr<<dfn[i]<<' ';
	}
	cerr<<endl;
	*/
	// segment tree
	build(1, 1, n);
	FOR(i, 1, k, 1){
		FOR(j, 0, szof(occ[i])-2, 1){
			int u = occ[i][j], v = occ[i][j+1]; // u, v = id. of nodes
			vector<pii> segs = Segs(u, v);
			for(pii p : segs){
				add_edge(1, p.F, p.S, dfn[u]);
			}
			add_edge(1, dfn[u], dfn[u], dfn[v]);
		}
	}
	
	// w
	FOR(i, 1, k, 1){
		int v = occ[i].front();
		w[leaf[dfn[v]]] = 1;
	}

	// SCC
	find_sccs();
	
	cout<<res - 1<<'\n';
	
	return 0;
	
}

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