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
題意：
n個節點，每個節點塗一個顏色，節點間部分有路徑形成樹，求所有路徑上所經過節點的不同顏色數目總和。
想法：
正向考慮時，要考慮每條路徑上存在的所有顏色數目，路徑多數目多，所以放棄正向思維，選擇逆向思維。
所有路徑的顏色總和⇌通過某顏色的路徑數總和
通過某顏色的路徑數總和=總路徑數-不經過此顏色的路徑數
→ 不經過此顏色的路徑數。
不經過此顏色的路徑數⇌樹中刪去此顏色構成的連通塊大小
不經過此顏色的路徑數=在除去顏色形成連通塊之外的路徑數+連通塊內形成的路徑數
→ 樹中刪去此顏色構成的連通塊大小
連通塊大小=子樹所有節點-以顏色爲根的子樹節點個數
注：代碼dfs函數已經求出連通塊內的路徑數，在主函數最後求的是連通塊之外的路徑數。
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>
#include<cmath>
#include<iostream>
using namespace std;
const int maxn=2e5+10;
vector<int>edge[maxn];//與結點i相連的結點 邊的數組
int siz[maxn];//以節點i爲根的子樹節點個數
int vis[maxn];
int col[maxn];
int sum[maxn];//以顏色i爲根的子樹且不包含此顏色的總和
typedef long long ll;
ll ans,n;
void dfs(int u,int fa)//（子節點，父節點）
{
siz[u]=1;
int pre=sum[col[u]];
int add=0;
for(int i=0;i<edge[u].size();++i)
{
int v=edge[u][i];//edge[u][i]表示與u相連的第i個子節點,v是整棵樹中的節點序號
if(v==fa)//找子節點
{//如果是父節點 就繼續
continue;
}
else
{
dfs(v,u);
}
siz[u]+=siz[v];
ll sonu=sum[col[u]]-pre; //在子樹v中，以顏色col[u]爲根的最高的子樹的大小
add+=sonu;//add記錄當前所有子樹中以col[u]爲根的節點總數
ll tmp=siz[v]-sonu;//子樹v中，除了u顏色之外的節點個數
ans-=tmp*(tmp-1)/2; //表示v子樹中，各個節點互相到達，無需經過col[u]
pre=sum[col[u]];//表示記錄上一次的sum
}
sum[col[u]]+=siz[u]-add;//sum[col[u]]加上以u爲根，屬於這個節點的連通塊的節點總數
}
int main()
{
int cas=1;
while(~scanf("%lld",&n))
{
memset(sum,0,sizeof sum);
memset(col,0,sizeof col);
memset(siz,0,sizeof siz);
memset(vis,0,sizeof vis);
int cnt=0;//顏色總種數
for(int i=1;i<=n;++i)
{
scanf("%d",&col[i]);
if(!vis[col[i]])
{
vis[col[i]]=1;
cnt++;
}
}
int u,v;
for(int i=1;i<=n;++i)
{
edge[i].clear();
}
for(int i=1;i<n;++i)
{
scanf("%d %d",&u,&v);
edge[u].push_back(v);//構成u v 之間的無向邊
edge[v].push_back(u);
}
ans=cnt*(n)*(n-1)/2;//所有顏色都存在的總路徑數
dfs(1,0);
for(int i=1;i<=n;++i)
{
if(!sum[i])
{
continue;
}
ll res=n-sum[i];
ans-=(res-1)*res/2;//此步驟減去的是 在i顏色節點的 上面部分，有n-sum[i]個節點，這些節點互通不需要經過i顏色節點。
}
printf("Case #%d: %lld\n",cas++,ans);
}
}
而標程其實是一樣逆向的思想，只是在求子樹中不含此顏色的節點數採取的是用dfs序來判斷的方式。
dfs序：將樹按dfs的方式排成一個序列，此點在序列中的序號稱爲dfs序。
①找出顏色中dfs序大於等於子根dfs序的位置
②若該位置指向該顏色的最後一個位置，代表已經是該子樹的最後
③若該位置大於該子根的最大的dfs序，代表已經不屬於這棵子樹
若不符合②③這兩種情況，就刪去這個位置，因爲這個位置含有以此位置爲根的子樹。
注：標程是在c++11才能運行的。
#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 200005;
int n , ca;
vector<int> e[N] , c[N];// e[i]與結點i相連的結點(邊數組), c[i]塗有顏色i有哪些點
int L[N] , R[N] , s[N] , f[N]; //L[i]結點i的dfs序, R[i]以i爲根結點的子樹中dfs序最大的點的dfs序, S[i]以i爲根節點的樹共有多少個結點, f[i]結點i的父親結點
void dfs(int x , int fa , int &&ncnt)
{
L[x] = ++ ncnt;
s[x] = 1 , f[x] = fa;
for (auto &y : e[x])
{
if (y != fa)
{
dfs(y , x , move(ncnt));
s[x] += s[y];
}
}
R[x] = ncnt;
}
bool cmp(const int& x , const int& y)
{
return L[x] < L[y];
}
void work()
{
for (int i = 0 ; i <= n ; ++ i) //從0開始 加了一個虛根 表示子樹路徑中不存在的顏色==
{
c[i].clear();
e[i].clear();
}
for (int i = 1 ; i <= n ; ++ i)
{
int x;
scanf("%d" , &x);
c[x].push_back(i);
}
for (int i = 1 ; i < n ; ++ i)
{
int x , y;
scanf("%d%d" , &x , &y);
e[x].push_back(y);
e[y].push_back(x);
}
e[0].push_back(1);
dfs(0 , 0 , 0);
LL res = (LL)n * n * (n - 1) / 2; //假設n個結點顏色均不同,每種顏色都貢獻了所有的路徑
for (int i = 1 ; i <= n ; ++ i)
{
if (c[i].empty()) //如果爲空,減一次(沒有這種顏色,減掉這種顏色所貢獻的路徑(全部))
{
res -= (LL)n * (n - 1) / 2;
continue;
}
c[i].push_back(0);
sort(c[i].begin() , c[i].end() , cmp);
for (auto &x : c[i])
{
for (auto &y : e[x])代表遍歷e[x]中所有的元素，每次遍歷時都把元素成爲y
{
if (y == f[x]) //若爲父親結點則continue(只考慮孩子結點)
continue;
int size = s[y];
int k = L[y];
while (1)
{
L[n + 1] = k;
auto it = lower_bound(c[i].begin() , c[i].end() , n + 1 , cmp); //找以孩子結點爲根節點的子樹中,相同顏色的結點
if (it == c[i].end() || L[*it] > R[y])
{
break;
}
size -= s[*it];
k = R[*it] + 1; //另一棵子樹
}
res -= (LL)size * (size - 1) / 2;
}
}
}
printf("Case #%d: %lld\n" , ++ ca , res);
}
int main()
{
while (~scanf("%d" , &n))
{
work();
}
return 0;
} ps：哇這個題真的是看了好久啊QAQ
看懂了才發現沒有那麼難，是自己不夠安靜去研究每一個變量的意思，其實自己靜下心來跑一遍就好了啊0.0