You are viewing a single comment's thread. Return to all comments →
Simple and concise solution encapsulating the ideas given by others:
#include <cmath> #include <cstdio> #include <vector> #include<stack> #include <iostream> #include <algorithm> using namespace std; long long int parent(long long int i, long long int* arr){ while(arr[i]>0) i=arr[i]; return i; } long long int bin(long long int a){ if(a<3) return 0; return ((a)*(a-1)*(a-2))/6; } int main() { long long int n; cin>>n; long long int parentarr[n+1]; for(long long int i=0;i<n+1;i++) parentarr[i]=-1; for(long long int i=1;i<n;i++){ long long int j,k; char col; cin>>j>>k>>col; long long int x=parent(j,parentarr); long long int y=parent(k,parentarr); if(x!=y&&col=='b'){ long long int c=parentarr[x]<=parentarr[y]?x:y; long long int d=parentarr[x]>parentarr[y]?x:y; parentarr[c]+=parentarr[d]; parentarr[d]=c; } } long long int total_triplets=bin(n); long long int to_subtract=0; for(long long int i=1;i<n+1;i++){ if(parentarr[i]>=-1) continue; long long int z=abs(parentarr[i]); to_subtract+=bin(z)+(n-z)*((z*(z-1))/2); } long long int sums=total_triplets-to_subtract; cout<<sums%(1000000007)<<endl; return 0; }
Seems like cookies are disabled on this browser, please enable them to open this website
Kundu and Tree
You are viewing a single comment's thread. Return to all comments →
Simple and concise solution encapsulating the ideas given by others: