标签:
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; }
HackerRank "Kundu and Tree" !!
标签:
原文地址:http://www.cnblogs.com/tonix/p/4996496.html