KunduAndTree.cpp

/*
https://www.hackerrank.com/contests/w5/challenges/kundu-and-tree
*/

#include <cstring>
#include <iostream>
#include <vector>
using namespace std;
 
#define MAX_N 100010
const int M = 1000000007;
int father[MAX_N]; // disjoint set
int num[MAX_N]; // number of elements in disjoint set
 
 
int find(int i)
{
    if (father[i] == i)
        return i;
    return father[i] = find(father[i]);
}
 
void merge(int i, int j)
{
    int x = find(i);
    int y = find(j);
    if (x == y) return;
    num[x] += num[y];
    father[y] = x;
}
 
int main()
{
    int n;
    cin >> n;
    for (int i = 0; i <= n; i++)
    {
        father[i] = i;
        num[i] = 1;
    }
    for (int i = 0; i < (n-1); i++)
    {
        int a, b;
        cin >> a >> b;
        char c;
        cin >> c;
        if (c == 'b')
        {
            merge(a, b);
        }
    }
    long long p1 = 0;
    long long p2 = 0;
    long long p3 = 0;
    for (int i = 1; i <= n; i++)
    {
        if (father[i] == i)
        {
            int x = num[i];
            p1 += x;
            p2 += x * x;
            p3 += x * x * x;
        }
    }
    long long result = (p1 * p1 * p1 - 3 * p2 * p1 + 2 * p3)/6;
    cout << result % M << endl;
}