#include <bits/stdc++.h>
#define int long long
#define P(x) do {if(debug) cout << x << endl;} while(false)
#define H(x) P(#x << ": " << x)
#define FR(i, a, b) for(int i = (a); i < (b); ++i)
#define F(i, n) FR(i, 0, n)
#define DR(i, a, b) for(int i = (b); i --> (a);)
#define D(i, n) DR(i, 0, n)
#define S(s) (int)(s).size()
#define ALL(x) (x).begin(), (x).end()
#define MI(x, a) (x) = min(x, a)
#define MA(x, a) (x) = max(x, a)
#define V vector
#define pb push_back
#define mp make_pair
using namespace std;
const bool debug = 0;
const int inf = 1e18;
int n,m;
vector<vector<int> > graph;
vector<int> color;
vector<int> par;
vector<bool> isdiam;
vector<int> maxdepth;
vector<int> subdiam;
vector<int> diam;
vector<bool> prohibited;
vector<int> half;
vector<int> halfbol;
vector<int> res;
int deepest=0; int deepestdepth=0;
int root;
int dfs1(int u, int p, int depth){
par[u] = p;
if(depth > deepestdepth){
deepestdepth = depth;
deepest = u;
}
maxdepth[u] = depth;
subdiam[u] = 1;
half[u] = u;
halfbol[u] = false;
int b = 0; int c = 0;
for(int i = 0; i<graph[u].size(); i++) if(graph[u][i]!=p){
int v = graph[u][i];
int a = dfs1(v,u,depth+1);
maxdepth[u] = max(maxdepth[u],a);
subdiam[u] = max(subdiam[u], subdiam[v]);
if(a>=b){
c = b;
b = a;
swap(graph[u][i], graph[u][0]);
half[u] = half[v];
if(halfbol[v]){
half[u] = par[half[u]];
}
halfbol[u] = !halfbol[v];
}
else if(a>c){
c = a;
}
}
subdiam[u] = max(subdiam[u], b+c);
return maxdepth[u];
}
vector<int> uniquec;
int cnt = 0;
vector<vector<int> > flagsubtree;
vector<vector<int> > flagall;
/*
process flagsubtree
set flag on subtrees if conditions are met
set flag on ancestor if conditions is met
dfs (process flagall in the meantime)
undo flagsubtree
*/
void dfs2(int u, int d){
H(u);
for(int c : flagsubtree[u]){
if(uniquec[c] == 0) cnt++;
uniquec[c]++;
}
for(int v : graph[u]) if(v!=par[u]){
if(maxdepth[v]+1 > subdiam[v]) {
flagsubtree[half[v]].push_back(color[u]);
P("set flag for color "<<color[u]<<" on "<<half[v]);
}
}
if(isdiam[u] && !prohibited[u] && u != root){
flagall[diam[(d-1)/2]].push_back(color[u]);
P("flagall for color "<<color[u]<<" on "<<diam[(d-1)/2]);
}
for(int v : graph[u]) if(v!=par[u]){
dfs2(v,d+1);
if(!flagall[u].empty()){
for(int c : flagall[u]){
if(uniquec[c]==0) cnt++;
uniquec[c]++;
}
flagall[u].clear();
}
}
res[u] = cnt;
for(int c : flagsubtree[u]){
if(uniquec[c] == 1) cnt--;
uniquec[c]--;
}
P("exit "<<u);
}
signed main() {
ios_base::sync_with_stdio(false);
cin.tie(0);
cin>>n>>m;
graph.resize(n);
color.resize(n);
F(i,n-1){
int a,b;
cin>>a>>b;
a--; b--;
graph[a].push_back(b);
graph[b].push_back(a);
}
F(i,n){
cin>>color[i];
color[i]--;
}
par.resize(n);
maxdepth.resize(n);
isdiam.resize(n,false);
subdiam.resize(n);
halfbol.resize(n);
half.resize(n);
dfs1(0,0,0);
root = deepest;
H(root);
deepestdepth = 0;
dfs1(root,-1,0);
P('a');
for(int u = deepest; u!=-1; u = par[u]){
isdiam[u] = true;
diam.push_back(u);
}
reverse(diam.begin(), diam.end());
//for(int i : diam){
// cout<<i<<" ";
//}
//cout<<endl;
prohibited.resize(n, false);
int j = -1;
for(int i = 0; i<diam.size(); i++){
int u = diam[i];
if(j>=i){
prohibited[diam[i]] = true;
}
for(int v : graph[u]) if(v!=par[u] && !isdiam[v]){
j = max(j, maxdepth[v]);
}
}
//for(bool b : prohibited){
// cout<<b<<" ";
//}
//cout<<endl;
P('b');
res.resize(n);
flagall.resize(n);
flagsubtree.resize(n);
uniquec.resize(n,0);
dfs2(root,0);
for(int i : res){
cout<<i<<endl;
}
return 0;
}
Compilation message
joi2019_ho_t5.cpp: In function 'long long int dfs1(long long int, long long int, long long int)':
joi2019_ho_t5.cpp:46:18: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for(int i = 0; i<graph[u].size(); i++) if(graph[u][i]!=p){
~^~~~~~~~~~~~~~~~
joi2019_ho_t5.cpp: In function 'int main()':
joi2019_ho_t5.cpp:163:18: warning: comparison between signed and unsigned integer expressions [-Wsign-compare]
for(int i = 0; i<diam.size(); i++){
~^~~~~~~~~~~~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
0 ms |
384 KB |
Output is correct |
| 2 |
Incorrect |
5 ms |
640 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
367 ms |
23416 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
544 ms |
32720 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
0 ms |
384 KB |
Output is correct |
| 2 |
Incorrect |
5 ms |
640 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
