CFR 734 E. Anton and Tree ( Ad hoc, Tree )

Problem - E - Codeforces

O( N )

```int N;
vi C;
vvi G;

void init(){
cin >> N;
C = vi( N );
for( int i = 0; i < N; ++i )
cin >> C[ i ];
G = vvi( N );
for( int i = 0; i < N - 1; ++i ){
int u, v; cin >> u >> v; --u, --v;
G[ u ].emplace_back( v );
G[ v ].emplace_back( u );
}
}

void solve(){
vi vis( N );
int far1;
queue< int > que;
que.emplace( 0 );
vis[ 0 ] = 1;
while( not que.empty() ){
int u = que.front();
que.pop();
far1 = u;
queue< int > q;
q.emplace( u );
while( not q.empty() ){
int v = q.front();
q.pop();
for( int x : G[ v ] ){
if( vis[ x ] ) continue;
vis[ x ] = 1;
if( C[ v ] == C[ x ] )
q.emplace( x );
else
que.emplace( x );
}
}
}
vis = vi( N );
int dmt;
queue< pii > d_u;
d_u.emplace( 0, far1 );
vis[ far1 ] = 1;
while( not d_u.empty() ){
int d, u; tie( d, u ) = d_u.front();
d_u.pop();
dmt = d;
queue< int > q;
q.emplace( u );
while( not q.empty() ){
int v = q.front();
q.pop();
for( int x : G[ v ] ){
if( vis[ x ] ) continue;
vis[ x ] = 1;
if( C[ v ] == C[ x ] )
q.emplace( x );
else
d_u.emplace( d + 1, x );
}
}
}
cout << ( dmt + 1 >> 1 ) << endl;
}
```