This submission is migrated from previous version of oj.uz, which used different machine for grading. This submission may have different result if resubmitted.
#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define ar array
const int mxN=2e5, MX=1e6;
int n, k, c[mxN], n2;
int d[mxN], s[mxN], p[mxN], hd[mxN], tin[mxN], tin2[mxN];
int who[mxN], cnt[mxN];
vector<int> oc[mxN], adj1[mxN], adj2[MX], adj3[MX], st, nodes;
bool vis[MX], out[MX];
void dfs1(int u=0) {
s[u]=1;
for (int& v : adj1[u]) {
adj1[v].erase(find(adj1[v].begin(), adj1[v].end(), u));
d[v]=d[u]+1, p[v]=u;
dfs1(v);
s[u]+=s[v];
if (s[v]>s[adj1[u][0]])
swap(v, adj1[u][0]);
}
}
void dfs2(int u=0, int h=0) {
static int timer=0;
hd[u]=h, tin[u]=timer++, tin2[tin[u]]=u;
if (adj1[u].empty())
return;
dfs2(adj1[u][0], h);
for (int i=1; i<adj1[u].size(); ++i)
dfs2(adj1[u][i], adj1[u][i]);
}
int lca(int u, int v) {
for (; hd[u]!=hd[v]; v=p[hd[v]])
if (d[hd[u]]>d[hd[v]])
swap(u, v);
return d[u]<d[v]?u:v;
}
void ae(int u, int v) {
adj2[u].push_back(v);
adj3[v].push_back(u);
}
void bld(int u=1, int l=0, int r=n-1) {
if (l==r) {
ae(u+k, c[tin2[l]]);
n2=max(n2, u+k+1);
return;
}
int mid=(l+r)/2;
ae(u+k, 2*u+k);
ae(u+k, 2*u+1+k);
bld(2*u, l, mid);
bld(2*u+1, mid+1, r);
}
void qry(int ql, int qr, int u=1, int l=0, int r=n-1) {
if (l>qr||r<ql)
return;
if (ql<=l&&r<=qr) {
nodes.push_back(u+k);
return;
}
int mid=(l+r)/2;
qry(ql, qr, 2*u, l, mid);
qry(ql, qr, 2*u+1, mid+1, r);
}
vector<int> gt(int u, int v) {
assert(hd[u]==hd[v]&&tin[u]<=tin[v]);
nodes.clear();
qry(tin[u], tin[v]);
return nodes;
}
void dfs3(int u) {
vis[u]=1;
for (int v : adj2[u])
if (!vis[v])
dfs3(v);
st.push_back(u);
}
void dfs4(int u, int rt) {
vis[u]=0, who[u]=rt;
cnt[rt]+=u<k;
for (int v : adj3[u])
if (vis[v])
dfs4(v, rt);
}
int main() {
ios::sync_with_stdio(0);
cin.tie(0);
cin >> n >> k;
for (int i=1; i<n; ++i) {
int u, v;
cin >> u >> v, --u, --v;
adj1[u].push_back(v);
adj1[v].push_back(u);
}
for (int i=0; i<n; ++i) {
cin >> c[i], --c[i];
oc[c[i]].push_back(i);
}
dfs1();
dfs2();
bld();
for (int i=0; i<n; ++i) {
ae(i+n2, c[i]);
if (i!=hd[i])
ae(i+n2, p[i]+n2);
}
for (int i=0; i<k; ++i) {
int l=oc[i][0];
for (int j=1; j<oc[i].size(); ++j)
l=lca(l, oc[i][j]);
for (int u : oc[i]) {
for (; hd[u]!=hd[l]; u=p[hd[u]])
ae(i, n2+u);
for (int j : gt(l, u))
ae(i, j);
}
}
n2+=n;
for (int i=0; i<n2; ++i)
if (!vis[i])
dfs3(i);
reverse(st.begin(), st.end());
for (int i : st)
if (vis[i])
dfs4(i, i);
for (int i : st)
for (int j : adj2[i])
if (who[i]!=who[j])
out[who[i]]=1;
int ans=69696969;
for (int i=0; i<k; ++i)
if (!out[who[i]])
ans=min(ans, cnt[who[i]]-1);
cout << ans;
return 0;
}
Compilation message (stderr)
capital_city.cpp: In function 'void dfs2(int, int)':
capital_city.cpp:32:17: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
32 | for (int i=1; i<adj1[u].size(); ++i)
| ~^~~~~~~~~~~~~~~
capital_city.cpp: In function 'int main()':
capital_city.cpp:120:18: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
120 | for (int j=1; j<oc[i].size(); ++j)
| ~^~~~~~~~~~~~~
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |
# | Verdict | Execution time | Memory | Grader output |
---|
Fetching results... |