HackerRank "Kundu and Tree" !!

Learnt from here: http://www.cnblogs.com/lautsie/p/3798165.html

Idea is: we union all pure black edges so we get 1+ pure black edge groups. Then we can simply pick only 1 vertex from each pure-black group to form a triangle. Then it will be all combination problem.

Only with some data type tuning to make it pass all tests:

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

unordered_map<int, int> parent;
unordered_map<int, long long> num;

#define M 1000000007

int find(int i)
{
    while(i != parent[i]) i = parent[i];
    return i;    
}
void merge(int i, int j) // merge j to i
{
    int pi = find(i);
    int pj = find(j);
    if (pi == pj) return;
        
    num[pi] += num[pj];
    parent[pj] = pi;
}

int main() 
{
    //    Get input
    int n; cin >> n;
    for (int i = 0; i <= n; i++)
    {
        parent[i] = i;
        num[i] = 1;
    }
    // union all black edges
    for (int i = 0; i < (n-1); i++)
    {
        int a, b;
        cin >> a >> b;
        char c;
        cin >> c;
        if (c == 'b')
        {
            merge(a, b);
        }
    }
    
    //    Now we have grouped all connected pure black edges
    
    //    Idea is, we pick 1 from 1 black set to make sure there's no pure
    //    black edges in one triangle. and we pick three
    long long p1 = 0;
    long long p2 = 0;
    long long p3 = 0;
    for(int i = 1 ; i <= n; i ++)
    {
        if(i == parent[i])
        {
            long long x = num[i];
            
            p1 += x; 
            p2 += (x * x); 
            p3 += (x * x * x);        
        }
    }
    
    long long ret = p1 * p1 * p1 - 3 * p2 * p1 + 2 * p3;
    cout << ((ret / 6) % M) << endl;
    
    return 0;
}

View Code