Hdu-6035 Colorful Tree(dfs)

Time Limit: 6000/3000 MS (Java/Others)    Memory Limit: 131072/131072 K (Java/Others) Total Submission(s): 375    Accepted Submission(s): 126 Problem Description There is a tree with  n  nodes, each of which has a type of color represented by an integer, where the color of node  i  is  ci . The path between each two different nodes is unique, of which we define the value as the number of different colors appearing in it. Calculate the sum of values of all paths on the tree that has  n(n−1)2  paths in total.   Input The input contains multiple test cases. For each test case, the first line contains one positive integers  n , indicating the number of node.  (2≤n≤200000) Next line contains  n  integers where the  i -th integer represents  ci , the color of node  i .  (1≤ci≤n) Each of the next  n−1  lines contains two positive integers  x,y   (1≤x,y≤n,x≠y) , meaning an edge between node  x  and node  y . It is guaranteed that these edges form a tree.   Output For each test case, output " Case # x :  y " in one line (without quotes), where  x  indicates the case number starting from  1  and  y  denotes the answer of corresponding case.   Sample Input 3 1 2 1 1 2 2 3 6 1 2 1 3 2 1 1 2 1 3 2 4 2 5 3 6   Sample Output Case #1: 6 Case #2: 29 分析：考虑每种颜色对答案的贡献,对于一种颜色我们考虑将其从图中扣去,那么只要计算剩下的每个连通快的大小就能算出这种 颜色对答案的贡献了,这一步我们把同颜色的点按dfs排序后可以o(n)的递归做. #include<bits/stdc++.h> #define N 200005 using namespace std; typedef long long ll; int n,x,y,cnt,Time,c[N],In[N],Out[N]; ll tot; vector<int> G[N],color[N]; void dfs(int u) { In[u] = ++cnt; for(int v : G[u]) if(!In[v]) dfs(v); Out[In[u]] = cnt; } ll got(ll x) { return x*(x-1)/2; } ll Count(int k,int s,int t,int l,int r) { int tot = r-l+1; ll ans = 0; while(s <= t) { int head = In[color[k][s]]+1,tail = Out[In[color[k][s]]]; tot -= tail - head + 2; s++; while(head <= tail) { int he = head,ta = Out[he]; int i = s; while(i <= t && In[color[k][i]] <= ta) i++; ans += Count(k,s,i-1,he,ta); s = i; head = ta + 1; } } return ans + got(tot); } struct cmp { bool operator() (int x,int y) { return In[x] <= In[y]; } }Cmp; int main() { cin.sync_with_stdio(false); while(cin>>n) { cnt = tot = 0; for(int i = 1;i <= n;i++) color[i].clear(),G[i].clear(),In[i] = Out[i] = 0; for(int i = 1;i <= n;i++) { cin>>c[i]; color[c[i]].push_back(i); } for(int i = 1;i < n;i++) { cin>>x>>y; G[x].push_back(y); G[y].push_back(x); } dfs(1); for(int i = 1;i <= n;i++) { sort(color[i].begin(),color[i].end(),Cmp); tot += n*(n-1ll)/2 - Count(i,0,color[i].size()-1,1,n); } cout<<"Case #"<<++Time<<": "<<tot<<endl; } }