Chinar - a symbol of Azerbaijan#12195

Chinar - a symbol of Azerbaijan (Анализа на решението)

If two vertices of the same color are connected by an edge, they can be merged into one — the answer will remain unchanged. Let's introduce the concept of a connected component: two vertices belong to the same connected component if they

Consider, for example, the tree on the left. On the right, the graph of its connected components is shown.

If we apply, for example, the operation , the component tree will be compressed:

The tree will be painted in a single color if and only if, after such painting operations with compression, exactly one component remains.

Let be the diameter of the component tree. The tree will be painted in a single color if and only if the diameter of the tree becomes , since a diameter of occurs only in a tree consisting of a single vertex.

We now find such a vertex that the length of the shortest path from it to any other vertex does not exceed (d + 1) / 2. Such a vertex always exists, because otherwise the diameter of the tree would be at least d + 1.

Finally, note that by applying (d + 1) / 2 painting operations to this vertex, we can paint the entire tree in a single color.

Algorithm implementation

Define the arrays.

vector<int> color, dist;
vector<vector<int> > g, g1;
vector<int> comp, used;

The function dfs1 performs a depth-first traversal of the tree and partitions its vertices into connected components. Each vertex is assigned a component number , which is stored in the array .

void dfs1(int v, int col, int id)
{
  if (used[v]) return;
  if (color[v] != col) return;

  used[v] = 1;
  comp[v] = id;


  for (int to : g[v])
    dfs1(to, col, id);
}

The function dfs2 performs a depth-first traversal of the tree and computes the distances from the source to each vertex. For a vertex , this distance is stored in dist[v]. The current distance from the source to vertex is denoted by .

void dfs2(int v, int d, int p = -1)
{
  dist[v] = d;
  for (int to : g1[v])
    if (to != p) dfs2(to, d + 1, v);
}

The main part of the program. Read the input data and initialize the arrays.

scanf("%d", &n);
color.resize(n + 1);
g.resize(n + 1);
comp.resize(n + 1);
used.assign(n + 1, 0);
for (i = 1; i <= n; i++)
  scanf("%d", &color[i]);

Read and construct an undirected tree.

for (i = 0; i < n - 1; i++)
{
  scanf("%d %d", &a, &b);
  g[a].push_back(b);
  g[b].push_back(a);
}

Partition the vertices of the tree into components. The total number of components is . The components are numbered starting from .

cnt = 0;
for (i = 1; i <= n; i++)
  if (!used[i])
  {
    cnt++;
    dfs1(i, color[i], cnt);
  }

Construct the graph — the component graph (the compressed tree).

g1.resize(cnt + 1);
for (v = 1; v <= n; v++)
  for (int to : g[v])
    if (comp[v] != comp[to])
      g1[comp[v]].push_back(comp[to]);

Compute the diameter of the component graph.

dist.resize(cnt + 1);

Start a depth-first search from vertex and fill the array , which stores the shortest distances from this vertex.

dfs2(1, 0, -1);

Vertex is the one located at the maximum distance from vertex .

a = max_element(dist.begin(), dist.begin() + cnt + 1) - dist.begin();

Start a depth-first search from vertex and fill the array again.

dfs2(a, 0, -1);

Vertex is the one farthest from . The distance between and is the diameter of the component graph.

b = max_element(dist.begin(), dist.begin() + cnt + 1) - dist.begin();

Print the answer.

printf("%d\n", (dist[b] + 1) / 2);