yaoxi-std 的博客

$\text{开}\mathop{\text{卷}}\limits^{ju\check{a}n}\text{有益}$

0%

P2664 树上游戏 题解

发表于 更新于 分类于

P2664 树上游戏 题解

题面

题目链接

解法

分别考虑每种颜色，对于颜色$col$，将所有$c_u=col$的$u$删去就得到若干连通块，颜色$col$对$sum_v$的贡献是$n-num_v$，$num_v$为$v$所在连通块大小。特别地，我们令$c_u=col$时$num_u=0$。

比较容易想到的是给每种颜色分别建立虚树计算贡献，时间复杂度为$O(n \log n)$。

但是这道题是有$O(n)$的做法滴。做法为两遍$dfs$，第一遍预处理出每个$col=c_{fa_u}$对应的$num_u$。$num_u$就是$siz_u$减去子树$u$中不经过其他颜色为$c_u$的节点直接与$u$连接的颜色为$c_u$的点的$siz$之和。$dfs$时顺便预处理出$col_i$表示根节点周围的颜色为$i$连通块大小。第二遍直接统计答案。

AC代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
/**
 * @file:           P2664.cpp
 * @author:         yaoxi-std
 * @url:            https://www.luogu.com.cn/problem/P2664
*/
// #pragma GCC optimize ("O2")
#include <bits/stdc++.h>
using namespace std;
#define int long long
#define resetIO(x) \
    freopen(#x ".in", "r", stdin), freopen(#x ".out", "w", stdout)
#define debug(fmt, ...) \
    fprintf(stderr, "[%s:%d] " fmt "\n", __FILE__, __LINE__, ##__VA_ARGS__)
template <class _Tp>
inline _Tp& read(_Tp &x) {
    bool sign = false;
    char ch = getchar();
    long double tmp = 1;
    for (; !isdigit(ch); ch = getchar())
        sign |= (ch == '-');
    for (x = 0; isdigit(ch); ch = getchar())
        x = x * 10 + (ch ^ 48);
    if (ch == '.')
        for (ch = getchar(); isdigit(ch); ch = getchar())
            tmp /= 10.0, x += tmp * (ch ^ 48);
    return sign ? (x = -x) : x;
}
template <class _Tp>
inline void write(_Tp x) {
    if (x < 0)
        putchar('-'), x = -x;
    if (x > 9)
        write(x / 10);
    putchar((x % 10) ^ 48);
}
const int MAXN = 1e5 + 10;
const int INFL = 0x3f3f3f3f3f3f3f3f;
int n, c[MAXN], son[MAXN], cnt[MAXN], col[MAXN], cur[MAXN];
int sum, len, siz[MAXN], num[MAXN], ans[MAXN];
vector<int> g[MAXN];
void dfs1(int u, int f) {
    int tmp = cur[c[u]];
    siz[u] = 1;
    cur[c[u]] = u;
    for (auto v : g[u]) {
        if (v == f)
            continue;
        son[u] = v;
        dfs1(v, u);
        siz[u] += siz[v];
    }
    cur[c[u]] = tmp;
    num[u] += siz[u];
    if (son[cur[c[u]]])
        num[son[cur[c[u]]]] -= siz[u];
    else
        col[c[u]] -= siz[u];
}
void dfs2(int u, int f) {
    if (son[cur[c[u]]])
        sum -= num[son[cur[c[u]]]];
    else
        sum -= col[c[u]];
    sum += num[u];
    ans[u] = n * len - sum;
    int tmp = cur[c[u]];
    cur[c[u]] = u;
    for (auto v : g[u]) {
        if (v == f)
            continue;
        son[u] = v;
        dfs2(v, u);
    }
    cur[c[u]] = tmp;
    sum -= num[u];
    if (son[cur[c[u]]])
        sum += num[son[cur[c[u]]]];
    else
        sum += col[c[u]];
}
signed main() {
    read(n);
    for (int i = 1; i <= n; ++i)
        read(c[i]), ++cnt[c[i]];
    for (int i = 1; i < MAXN; ++i)
        if (cnt[i])
            ++len, col[i] = n;
    for (int i = 1; i < n; ++i) {
        int u, v;
        read(u), read(v);
        g[u].push_back(v);
        g[v].push_back(u);
    }
    dfs1(1, 0);
    for (int i = 1; i < MAXN; ++i)
        sum += col[i];
    num[1] = 0;
    dfs2(1, 0);
    for (int i = 1; i <= n; ++i)
        write(ans[i]), putchar('\n');
    return 0;
}