Before the Beginning
转载请将本段放在文章开头显眼处，如有二次创作请标明。
原文链接：https://www.codein.icu/lp5588/
题意
题目给的描述足够清晰，照搬过来：
解法
由特殊到一般，逐步分析。
对每个颜色进行讨论。
颜色未出现
颜色的值域为 $[1,n]$，一共 $n$ 个点，只要有点的颜色重复，就会有颜色未在树上出现。
由于未在树上出现，任意点对都是满足条件的点对。
每个点与剩余 $n-1$ 个点构成点对，除去重复，答案即 $\frac{n \times (n – 1)}{2}$。
颜色出现一次
颜色出现一次时，经过这个点的路径即合法。
将该颜色的点当成树根，旋转树。
经过树根的点对就是合法的。
对于每一个儿子为根的子树内的点，都可以与非该子树内的点组成点对。
考虑如何不重不漏地统计答案。
定义 $size_i$ 为以 $i$ 为根的子树的大小。
每遍历到一个儿子时，统计它与已遍历儿子的点对。
for(int p = head[u];p;p=E[p].next)
{
v = E[p].to;
if(v == fa)
continue;
dfs(v,u);
ans[u] += size[u] * size[v];
size[u] += size[v];
}
ans[u] += (n - size[u]) * size[u];特别地，将父亲也看作换根后的儿子看待，在最后特殊处理。
颜色在链上出现
可以发现，链有两个端点。
旋转树，将端点当做根，同颜色的节点一定只在一个儿子的子树中出现。
合法的路径为：
从端点和未出现同色节点的子树中的节点 到
另一个端点，和另一个端点为根时未出现同色节点的子树中的节点。
这部分的大小，用节点总数减去同色节点出现的子树大小得出。
现在的问题就是如何寻找端点。
根据性质，可以发现：
将树旋转为以端点为根的形态，有且只有一个儿子为根的子树中会有同色节点。
记录有多少儿子子树出现了同色节点，如果唯一则为端点。
特别地，父亲节点要当做儿子特殊处理。
记 $cnt_i$ 为搜索过程中颜色 $i$ 出现的次数。
向下搜索，回溯到源点，若颜色 $i$ 出现的次数变化了，说明该子树中有同色节点出现。
将该儿子记录下来。
int has = 0,hasSon;
int beginColorNum = cnt[w[u]];
for(int p = head[u];p;p=E[p].next)
{
v = E[p].to;
if(v == fa)
continue;
int temp = cnt[w[u]];
dfs(v,u);
if(cnt[w[u]] != temp)
++has,hasSon = v;
}
//处理父亲节点，搜索完儿子子树后若记录数量仍不足，说明父亲子树中也有同色点
if(beginColorNum != 0 || cnt[w[u]] + 1 != num[w[u]])
++has;
++cnt[w[u]];
if(has == 1)
{
//dosomething...
}这样就可以找出端点了。
一条链有且只有两个端点，如果出现多于两个端点，说明链在某处分叉了，一定没有符合条件的路径，答案为 $0$。
记录端点数量。
if(has == 1)
{
int t;//计算合法的起点的数量
if(hasSon)
t = n - size[hasSon];//链在儿子为根的子树中，减去该子树，剩余节点都可为起点
else
t = size[u];//链在父亲为根的子树中，当前子树都可为起点。
if(pointNum[w[u]] == 0)//第一个端点
firstPos[w[u]] = u,preSize[w[u]] = t;
else if(pointNum[w[u]] == 1)//两个端点都已经确定
{
//答案数量为合法起点、终点相乘
f2[w[u]] = (long long)t * preSize[w[u]];
}
pointNum[w[u]]++;
}颜色不在链上出现
一条链有且只有两个端点，如果出现多于两个端点，说明链在某处分叉了，一定没有符合条件的路径，答案为 $0$。
代码
将上述的做法综合起来。十年OI一场空，不开…
#include <cstdio>
template <typename T>
inline void read(T &r)
{
static char c;
r = 0;
for (c = getchar(); c > '9' || c < '0'; c = getchar());
for (; c >= '0' && c <= '9'; r = (r << 1) + (r << 3) + (c ^ 48), c = getchar());
}
const int maxn = 1e6 + 10;
struct node
{
int to, next;
} E[maxn << 1];
int head[maxn];
inline void add(const int &x, const int &y)
{
static int tot = 0;
E[++tot].to = y;
E[tot].next = head[x];
head[x] = tot;
}
int n;
int w[maxn];
int colorNum[maxn], size[maxn];
long long f1[maxn]; //颜色只在当前点出现时的答案
long long f2[maxn]; //f2[i] 颜色i在链上时的答案
int cnt[maxn]; //cnt[i] 颜色i在当前出现的次数
int pointNum[maxn], preSize[maxn];
void dfs(int u, int fa)
{
size[u] = 1;
int v;
int has = 0, hasSon = 0;
int beginColorNum = cnt[w[u]];
for (int p = head[u]; p; p = E[p].next)
{
v = E[p].to;
if (v == fa)
continue;
int temp = cnt[w[u]];
dfs(v, u);
f1[u] += (long long)size[u] * size[v];
size[u] += size[v];
if (cnt[w[u]] != temp)
++has, hasSon = v;
}
f1[u] += (long long)(n - size[u]) * size[u]; //父亲节点
if (beginColorNum != 0 || cnt[w[u]] + 1 != colorNum[w[u]])
++has;
++cnt[w[u]];
if (has == 1)
{
int t = hasSon ? n - size[hasSon] : size[u];
if (pointNum[w[u]] == 0)
preSize[w[u]] = t;//事实上只需要记录大小就够了，不用确定点
else if (pointNum[w[u]] == 1)
f2[w[u]] = (long long)t * preSize[w[u]];
pointNum[w[u]]++;
}
}
int lastPos[maxn]; //记录某个颜色最后出现的点，用于处理只出现一次的情况
int main()
{
read(n);
for (int i = 1; i <= n; ++i)
read(w[i]), colorNum[w[i]]++, lastPos[w[i]] = i;
int x, y;
for (int i = 1; i < n; ++i)
{
read(x);
read(y);
add(x, y);
add(y, x);
}
dfs(1, 0);
for (int i = 1; i <= n; ++i)
{
if (colorNum[i] == 0)
printf("%lld\n", (long long)n * (n - 1) >> 1);
else if (colorNum[i] == 1)
printf("%lld\n", f1[lastPos[i]]);
else if (pointNum[i] == 2)
printf("%lld\n", f2[i]);
else
puts("0");
}
return 0;
}笔者认为这道题有蓝题难度了，当然可能是笔者太蒟了……