*[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;

}