Hackerrank--Kundu and Tree

タイトルリンク
Kundu is true tree lover. Tree is a connected graph having N vertices and N-1 edges. Today when he got a tree, he colored each edge with one of either red( r ) or black( b ) color. He is interested in knowing how many triplets(a,b,c) of vertices are there , such that, there is atleast one edge having red color on all the three paths i.e. from vertex a tob, vertex b to c and vertex c to a . Note that (a,b,c), (b,a,c) and all such permutations will be considered as the same triplet.
If the answer is greater than 109 + 7, print the answer modulo (%) 109 + 7.
Input FormatThe first line contains an integer N, i.e., the number of vertices in tree. The next N-1 lines represent edges: 2 space separated integers denoting an edge followed by a color of the edge. A color of an edge is denoted by a small letter of English alphabet, and it can be either red( r ) or black( b ).
Output FormatPrint a single number i.e. the number of triplets.
Constraints1 ≤ N ≤ 105A node is numbered between 1 to N.
Sample Input

5 1 2 b 2 3 r 3 4 r 4 5 b 

Sample Output

4 

Explanation
Given tree is something like this.
(2,3,4) is one such triplet because on all paths i.e 2 to 3, 3 to 4 and 2 to 4 there is atleast one edge having red color.(2,3,5), (1,3,4) and (1,3,5) are other such triplets. Note that (1,2,3) is NOT a triplet, because the path from 1 to 2 does not have an edge with red color.
大体の問題:1本の木を与えて、2つの木の辺があって、1つは赤('r')、もう1つは黒('b')で、木の中で3つの点を見つけて、2つの点の間の経路の上ですべて赤い辺があることができなくて、
何種類の異なる探し方があるかを聞く.(a,b,c,b a cは1種のみ)
構想:要求を満たす3つの点a,b,cがあれば,aからbの中の赤い辺を取り除くと,aとbは連通せず,同じ理屈で,aからb,bからc,aからcに
の赤い辺がすべて取り除かれると、a,b,cは異なる連通成分に属します.これにより、k個のスタックがあり、各スタックにa[i]個があり、このk個のスタックから3個のポイントを選択し、3個のポイントを満たすように変換される.
両両は同じ山にいないので,どれだけの中取法があるかを聞く.そのため、次の方法があります.

1 long long sum = 0; 2 for (int i = 1; i <= k; i++) 3     for (int j = i + 1; j <= k; j++) 4         for (int t = j + 1; t <= k; t++) 5             sum += a[i] * a[j] * a[t];

しかし複雑度はO(n^3)であり,明らかに問題要件を満たすことができない.O(n)の複雑さに最適化しなければならない.think about it, how to optimize the solution....???
Accepted Code:

 1 #include <iostream>

 2 #include <cstring>

 3 #include <cstdlib>

 4 #include <vector>

 5 using namespace std;  6 

 7 const int MOD = 1000000000 + 7;  8 const int MAX_N = 100005;  9 typedef long long LL; 10 vector<int> G[MAX_N]; 11 int N, cmp[MAX_N]; 12 LL cnt[MAX_N], A[MAX_N], B[MAX_N], C[MAX_N]; 13 bool vis[MAX_N]; 14 

15 void dfs(int u, int k) { 16     cmp[u] = k; vis[u] = true; 17     for (int i = 0; i < G[u].size(); i++) { 18         int v = G[u][i]; 19         if (!vis[v]) dfs(v, k); 20  } 21 } 22 int main(void) { 23      while (cin >> N) { 24          for (int i = 1; i <= N; i++) G[i].clear(); 25          for (int i = 0; i < N; i++) { 26              int a, b; char c; cin >> a >> b >> c; 27              if (c != 'r') G[a].push_back(b), G[b].push_back(a); 28  } 29         memset(vis, false, sizeof(vis)); 30         int k = 1; //       31         for (int i = 1; i <= N; i++) if (!vis[i]) dfs(i, k++); 32         k--; 33         memset(cnt, 0, sizeof(cnt));　//           34         for (int i = 1; i <= N; i++) cnt[cmp[i]]++; 35         A[k] = cnt[k];　//A[i] = cnt[i] + cnt[i + 1] + ... + cnt[k] 36         for (int i = k - 1; i >= 3; i--) A[i] = (A[i + 1] + cnt[i]) % MOD; 37         for (int i = 2; i < k; i++) B[i] = (cnt[i] * A[i + 1]) % MOD; //B[i] = cnt[i] * A[i + 1]. 38         C[k - 1] = B[k - 1]; //C[i] = B[i] + B[i + 1] + ... + B[k - 1] 39         for (int i = k - 2; i >= 2; i--) C[i] = (C[i + 1] + B[i]) % MOD; 40         LL sum = 0; 41         for (int i = 1; i <= k - 2; i++) sum = (sum + cnt[i] * C[i + 1]) % MOD; 42         cout << sum << endl; 43  } 44     return 0; 45 }