Topcoder SRM 640 Div. 1 250 - ChristmasTreeDecoration

Problem Statement:
SRM 640 Div. 1 250 - ChristmasTreeDecoration

Solution:
We can employ a greedy scheme:
1. For each possible remaining edge, choose the one that connects two vertices with different colors and that does not result in a loop.
2. If it is not possible anymore to do that, the remaining edges needed to form the tree will all need to connect vertices of the same color.



Why can we do this? I think it is possible to prove this by exchange argument or some by contradiction. Notice that if we choose all edges that end with vertices of different color, we will have a graph. If it forms a tree then everything is good. Otherwise, there are some loops in the graph, we can remove any edges from the graph until we form a tree. Notice all these vertices will still be connected. For remaining nodes that are not yet connected to this component, we add the required edges, which by initial assumption, will be those edges ending with pairs of vertices of the same color.

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <vector>
using namespace std;

int par[100];
int mark[203];
void init_dsu() {
 for(int i=0;i<100;++i) par[i] = i;
}

int find_par(int i) {
 return (par[i] == i ? i : (par[i] = find_par(par[i])));
}

void merge(int i, int j){
 int x = find_par(i);
 int y = find_par(j);
 par[x] = y;
}

class ChristmasTreeDecoration {
public:
 int solve(vector<int> col, vector<int> x, vector<int> y) {
  init_dsu();
  int N = col.size();
  int M = x.size();
  int chosen = 0;
  for(int i=0;i<M;++i){
   mark[i] = 0;
  }
  for(int i=0;i<N-1;++i){
   bool found = false;
   for(int i=0;i<M;++i){
    if(found) break;
    if(mark[i]) continue;
    if(col[x[i]-1] == col[y[i]-1]) continue;
    if(find_par(x[i]-1) == find_par(y[i]-1)) continue;
    mark[i] = 1;
    merge(x[i]-1, y[i]-1);
    ++chosen;
    found = true;
   }
   if(!found) break;
  }
  return N-1 - chosen;
 }
};