洛谷 P11251 [GESP202409 八级] 美丽路径-普及/提高-

题目描述

小杨有一棵包含 $n$ 个节点的树，节点从 $1$ 到 $n$ 编号。每个节点要么是白色，要么是黑色。

对于树上的一条简单路径（不经过重复节点的路径），小杨认为它是美丽的当且仅当路径上相邻节点的颜色均不相同。例如下图，其中节点 $1$ 和节点 $4$ 是黑色，其余节点是白色，路径 $2-1-3-4$ 是美丽路径，而路径 $2-1-3-5$ 不是美丽路径（相邻节点 $3$ 和 $5$ 颜色相同）。

对于树上一条简单路径，小杨认为它的长度是路径包含的节点数量。小杨想知道最长美丽路径的长度是多少。

输入格式

第一行包含一个正整数 $n$，表示节点数量。
第二行包含 $n$ 个整数 $c_1,c_2,c_3,\dots c_n$，代表每个节点的颜色。如果 $c_i=0$，节点 $i$ 为白色；如果 $c_i=1$，节点 $i$ 为黑色。
之后 $n-1$ 行，每行两个正整数 $u_i,v_i$，表示树上一条边。

输出格式

输出一行一个整数表示答案。

输入输出样例 #1

输入 #1

5
1 0 0 1 0
1 2
3 5
4 3
1 3

输出 #1

4

输入输出样例 #2

输入 #2

5
0 0 0 0 0
1 2
2 3
3 4
4 5

输出 #2

1

说明/提示

子任务	占比	$n$	特殊约定
$1$	$30\%$	$\le 1000$	树的形态是一条链
$2$	$30\%$	$\le 1000$	无
$3$	$40\%$	$\le 10^5$	无

对全部的测试数据，保证 $1\le n\le 10^5$，$0\le c_i\le 1$，保证给出的是一棵树。

solution

对于每一个节点 u 统计每一条根节点和本节点不一致的子树的交替路径长度，计算最长和次长长度和，并向父节点返回最长长度

代码

#include 
#include "bit"
#include "vector"
#include "unordered_set"
#include "set"
#include "queue"
#include "stack"
#include "algorithm"
#include "bitset"
#include "cstring"

using namespace std;

const int N = 1e5 + 5;
int n, a[N], ans;
vector<int> e[N];

int dfs(int u, int p) {
    int l = 0, s = 0;
    for (int v: e[u]) {
        if (v != p) {
            if (a[v] != a[u]) {
                int m = dfs(v, u);
                if (m > l) s = l, l = m;
                else s = m;
            }else{
                dfs(v, u);
            }
        }
    }
    ans = max(ans, l + s + 1);
    return l + 1;
}

int main() {
    cin >> n;

    for (int i = 1; i <= n; i++) cin >> a[i];
    for (int i = 1; i < n; i++) {
        int x, y;
        cin >> x >> y;
        e[x].push_back(y);
        e[y].push_back(x);
    }

    dfs(1, 0);
    cout << ans;
}

结果