*[hackerrank]Kundu and Tree

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

此题是有一棵树，树中的边有红有黑。求节点(a,b,c)的组合个数，使得ab,bc,ac都至少经过一条红边。

思路是：只有红边有意义，那么黑边连接的点都可以归为等价类（并查集），之后就是组合问题。记X1+X2+X3+...+Xn为P1，二次方的和为P2，三次方的和为P3。那么结果为(P1^3-3*P2*P1-2*P3)/6。6是三个不相同元素的排列情况，P1^3是三个元素所有取的排列次数，中间是两个在一类的情况，*3表示有三种排列情况，最后是如果三个元素在一类，那么会在第二种情况中计算两遍。

还有个更简单的算法就是C(n,3)，减所有C(x,3)，再减C(x,2)*(n-x)表示三个点在一起和两个点在一起的。

#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;
}