2-牛客周赛 Round 138题解

链接：https://ac.nowcoder.com/acm/contest/131111/D 来源：牛客网
小苯有一棵 $n$ 个节点的有根树，根为 $1$ 。每个节点 $i$ 有一个颜色 $c_i$ 。
小苯定义一条路径 $u \to v, u\neq v$ （ $u$ 是 $v$ 的祖先）是“同色路径”，当且仅当路径上所有节点的颜色相同
注意：路径必须是祖先 $\rightarrow$ 后代关系，且方向必须从祖先到后代。
小苯想知道，树中一共有多少条不同的同色路径。
你的任务是求出同色路径的数量。
输入描述:#
每个测试文件均包含多组测试数据。
第一行输入一个整数 $T\ (1 \le T \le 10^5)$ ，代表数据组数，每组测试数据描述如下：
第一行包含一个整数 $n\ (1 \le n \le 2 \times 10^5)$ 。
第二行包含 $n$ 个整数 $c_1, c_2, \dots, c_n\ (1 \le c_i \le 10^9)$ 。
接下来 $n - 1$ 行，每行包含两个整数 $u_i, v_i\ (1 \le u_i, v_i \le n)$ ，表示树中有一条边 $(u_i, v_i)$ ，保证构成一棵以 $1$ 为根的树。
除此之外，保证单个测试文件的 $\sum n \le 2 \times 10^5$ 。
输出描述:#
对于每组测试数据：输出一个整数，表示同色路径的数量。
示例1#
输入#
1
2
2
5
3
1 1 2 1 2
4
1 2
5
2 3
6
2 4
7
4 5
8
6
9
1 2 2 3 3 1
10
1 2
11
2 3
12
3 4
13
4 5
14
5 6
输出#
1
3
2
2
说明#
对于第一组测试数据，合法路径有： $\{1 \rightarrow 2,\ 2 \rightarrow 4,\ 1 \rightarrow 2 \rightarrow 4\}$ ，共 $3$ 个。

对每个结点记录以该结点为结尾的最长路径长度，应为自顶向下的树形DP

1
#include <bits/stdc++.h>
2

3
using namespace std;
4

5
const int N = 2e5 + 1;
6

7
int c[N];
8

9
void solve() {
10
    int n;cin>>n;
11
    for (int i = 0; i < n; ++i) {
12
        cin >> c[i];
13
    }
14

15
    vector<vector<int>> g(n);
16
    vector<int> fa(n);
17
    vector<int> dp(n);
18
    for (int i = 1; i < n; ++i) {
19
        int u, v;
20
        cin >> u >> v;--u;--v;
21
        g[u].push_back(v);
22
        g[v].push_back(u);
23
    }
24

25
    queue<int> q;
26
    q.push(0);
27
    dp[0] = 1;
28
    int ans = 0;
29
    while (q.size()) {
30
        int sz = q.size();
31
        while (sz--) {
32
            int t = q.front();
33
            q.pop();
34

35
            for (auto& son : g[t]) {
36
                if (son != fa[t]) {
37
                    fa[son] = t;
38
                    if (c[son] == c[t]) {
39
                        dp[son] = dp[t] + 1;
40
                        ans += dp[t];
41
                    } else {
42
                        dp[son] = 1;
43
                    }
44
                    q.push(son);
45
                }
46
            }
47
        }
48
    }
49

50
    cout << ans << '\n';
51
}
52

53
int main() {
54
    ios::sync_with_stdio(false);cin.tie(nullptr);
55
    int t; cin>>t;
56
    while (t--) {
57
        solve();
58
    }
59
}

音乐

音乐

输入描述:#

输出描述:#

示例1#

输入#

输出#

说明#

文章分享

评论区

音乐

目录