解題報告：HDU_6035 Colorful Tree 樹上計數

題目鏈接

題意：

給定一棵樹，樹上每個結點都有一種顏色，詢問路上任意兩點之間的路徑上包含的不同顏色數目和。

思路：
正着求不好求，先認爲所有的路徑上包含所有出現過的顏色，那麼只需要減去每個顏色沒有出現過的路徑之和。

官方題解：

單獨考慮每一種顏色，答案就是對於每種顏色至少經過一次這種的路徑條數之和。反過來思考只需要求有多少條路徑沒有經過這種顏色即可。直接做可以採用虛樹的思想（不用真正建出來），對每種顏色的點按照 dfs 序列排個序，就能求出這些點把原來的樹劃分成的塊的大小。這個過程實際上可以直接一次 dfs 求出。

代碼：

#include<bits/stdc++.h>

const int N = 2e5+10;
using namespace std;

int n;
vector<int>G[N];
int col[N],num[N],son[N];
bool used[N];
long long ans ;

void dfs(int x,int fa = 0){
   son[x] = 1;
   int& c = col[x];
   int last = num[c] , all = 0;
   for(int i=0;i<G[x].size();i++){
      int j = G[x][i];
      if(j!=fa){
         dfs(j,x);
         son[x] += son[j];
         all += num[c] - last;
         int cnt = son[j] - num[c] + last;
         last = num[c];
         ans -= 1LL * cnt * (cnt-1) / 2;
      }
   }num[c] += son[x] - all ;
}


int main()
{
   int cas = 0;
   while(scanf("%d",&n)==1){
      ans = 0;
      memset(used,0,sizeof(used));
      for(int i=1;i<=n;i++){
         G[i].clear();
         scanf("%d",&col[i]);
         num[col[i]] = 0;
         used[col[i]] = 1;
      }for(int i=1,s,e;i<n;i++){
         scanf("%d%d",&s,&e);
         G[s].emplace_back(e);
         G[e].emplace_back(s);
      }dfs(1);int cnt = 1;
      for(int i=1;i<=n;i++)if(used[i]&&i!=col[1]){
         cnt ++;
         int tmp = n - num[i];
         ans -= 1LL * tmp * (tmp-1) / 2;
      }ans += 1LL * n * (n-1) / 2 * cnt ;
      printf("Case #%d: %I64d\n",++cas,ans);
   }return 0;
}