P5593 题解

题目

分析

首先考虑什么样的颜色能被链覆盖。

容易想到当某种颜色恰巧在一条链上会被覆盖。

所以只需要判断一种颜色是否能构成链即可，链的贡献也很好计算。

算法

考虑链的性质：有且仅有两个端点。

凭借这个性质，可以判断一种颜色是否在一条链上。

在 dfs 中考虑各种情况。

假设一个点 \(u\)，其颜色为 \(c\)，有以下 \(2\) 种情况判断 \(u\) 是否为端点。

1. 如果 \(u\) 的子树中没有颜色 \(c\) 的节点，那么 \(u\) 点是一个端点。
2. 满足两个条件。
   1. 除了 \(u\) 的子树没有颜色 \(c\) 的节点。
   2. \(u\) 的所有儿子中，只有一个儿子的子树中有颜色 \(c\)。

判断端点数量是否等于 \(2\)，等于 \(2\) 说明该种颜色恰好构成一条链。

还有一个特殊情况：某个颜色可能只有一个点，需要直接算答案。

如果是一条链的话，假设两条节点为 \(p, q\)，对答案的贡献有两种情况。

1. 如果 \(p, q\) 均是满足上述第 \(1\) 种情况的，答案为 \(sz_p \times sz_q\) 。
2. 如果 \(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_29
    freopen("read.in", "r", stdin);
    freopen("out.out", "w", stdout);
#endif

    RD(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;
}

七月流火，九月授衣。