题目
分析
首先考虑什么样的颜色能被链覆盖。
容易想到当某种颜色恰巧在一条链上会被覆盖。
所以只需要判断一种颜色是否能构成链即可,链的贡献也很好计算。
算法
考虑链的性质:有且仅有两个端点。
凭借这个性质,可以判断一种颜色是否在一条链上。
在 dfs 中考虑各种情况。
假设一个点 \(u\),其颜色为 \(c\),有以下 \(2\) 种情况判断 \(u\) 是否为端点。
- 如果 \(u\) 的子树中没有颜色 \(c\) 的节点,那么 \(u\) 点是一个端点。
- 满足两个条件。
- 除了 \(u\) 的子树没有颜色 \(c\) 的节点。
- \(u\) 的所有儿子中,只有一个儿子的子树中有颜色 \(c\)。
判断端点数量是否等于 \(2\),等于 \(2\) 说明该种颜色恰好构成一条链。
还有一个特殊情况:某个颜色可能只有一个点,需要直接算答案。
如果是一条链的话,假设两条节点为 \(p, q\),对答案的贡献有两种情况。
- 如果 \(p, q\) 均是满足上述第 \(1\) 种情况的,答案为 \(sz_p \times sz_q\) 。
- 如果 \(p\) 是满足上述第 \(2\) 种情况的,考虑一下,因为 \(p\) 的所有儿子中,只有一个儿子的子树中有颜色 \(c\),假设该儿子为 \(pos\),所以 \(p\) 的其他儿子的子树中的节点也能对答案有贡献,故答案为 \((n - sz_{pos}) \times sz_q\)。
如何实现见代码。
#ifdef ONLINE_JUDGE
#else
#define FRZ_29
#endif#include <iostream>
#include <cstdio>
typedef long long LL;using namespace std;void RD() {}
template<typename T, typename... U> void RD(T&x, U&... arg) {x = 0; int f = 1; char ch = getchar();while (ch > '9' || ch < '0') { if (ch == '-') ch = getchar(); ch = getchar(); }while (ch <= '9' && ch >= '0') x = (x << 3) + (x << 1) + ch - '0', ch = getchar();x *= f; RD(arg...);
}void WT() {}
template<typename T> void WT(T x) {if (x < 0) { putchar('-'); x = -x; }if (x > 9) WT(x / 10);putchar(x % 10 + '0');
}const int N = 3e6 + 5;#define LF(i, __l, __r) for (int i = __l; i <= __r; i++)
#define RF(i, __r, __l) for (int i = __r; i >= __l; i--)int head[N], ver[N << 1], Next[N << 1], _tot = 1;
int n, a[N], tot[N], nos[N], enos[N], cnt[N], sz[N];
LL ans, ans1[N], ans2[N];void add(int u, int v) {ver[++_tot] = v;Next[_tot] = head[u], head[u] = _tot;
}/*
如果 flag == 1,才能判断为端点
这是由 flag 的增加方式决定的
*/
void dfs(int u, int _f) {int c = a[u], flag = 0, pos = 0;int k = cnt[c];sz[u] = 1;for (int i = head[u]; i; i = Next[i]) {int v = ver[i];if (v == _f) continue;int la = cnt[c];dfs(v, u);ans1[u] += 1LL * sz[u] *sz[v];sz[u] = sz[u] + sz[v];if (la != cnt[c]) flag++, pos = v; // 与判断是否为端点的条件2相关} // 如果该条件被满足两次pos自然就无用了ans1[u] += 1LL * sz[u] * (n - sz[u]);if (k || cnt[c] != tot[c] - 1) flag++; // 与判断是否为端点的条件1相关cnt[c]++;if (flag == 1) {if (!enos[c]) nos[c] = u;else {int p = pos ? n - sz[pos] : sz[u];ans2[c] = 1LL * p * sz[nos[c]];}enos[c]++;}
}int main() {
#ifdef FRZ_29freopen("read.in", "r", stdin);freopen("out.out", "w", stdout);
#endifRD(n);LF(i, 1, n) {RD(a[i]); tot[a[i]]++; nos[a[i]] = i;}LF(i, 1, n - 1) {int u, v; RD(u, v);add(u, v), add(v, u);}dfs(1, 0);LF(i, 1, n) {if (tot[i] == 0) ans ^= 1LL * n * (n - 1) / 2;else if (tot[i] == 1) ans ^= ans1[nos[i]];else if (enos[i] == 2) ans ^= ans2[i];else ans ^= 0;}WT(ans);return 0;
}
七月流火,九月授衣。