INVENT spoj.com Solution

Problem : http://www.spoj.com/problems/INVENT/

Algorithm : Union-Find

1. Sort the given edges according to increasing order of the weight.
2. Use Union-Find to build an MST with edge weights incremented by 1. ie. weight1 = weight+1.
3. Subtract the cost equal to the number of edges to get a unique MST from the MST we just made.

Source Ref - https://github.com/VitSalis/SPOJ/blob/master/INVENT/invent.cpp

Source:

#include <stdio.h> 
#include <algorithm>

using namespace std; 

int ranks[15001];
int num[15001];
int P[15001];

struct edge{
 int from, to, weight;
};

bool operator <(edge a, edge b){
 return a.weight < b.weight;
}

void makeSet(int x){
 P[x] = x;
 ranks[x] = 1;
 num[x] = 1;
}

int findSet(int x){  
 while (P[x]!=x){ 
  x = P[x]; 
 }   
 return x;
}

void uni(int x, int y){
 x = findSet(x);
 y = findSet(y);
 if (ranks[x] > ranks[y]){
  P[y] = x;
  num[x] += num[y];
 }
 else{
  P[x] = y;
  num[y] += num[x];
  if (ranks[x] == ranks[y]){
   ranks[y] += 1;
  }
 }
}

int main(){
 int nCases;
 scanf("%d", &nCases);

 while (nCases--){
  int V;
  scanf(" %d", &V);
  
  int E = V - 1;
  
  edge array[15001];
  unsigned long long cost = 0;

  for (int i = 0; i < E; ++i){ 
   scanf("%d %d %d", &array[i].from, &array[i].to, &array[i].weight);
   makeSet(array[i].from);
   makeSet(array[i].to);
  }

  sort(array, array + E);
  for (int i = 0; i < E; ++i){
   int u = findSet(array[i].from);
   int v = findSet(array[i].to);  
   cost += ((unsigned long long)num[u] * num[v] * (array[i].weight + 1));
   uni(array[i].from, array[i].to); 
  }

  printf("%lld\n", cost - E);
 }
}