大二下课业比较重,只能挤出时间来训练
周三
E. Anton and Tree(并查集+树的直径)
这题我下意识以为是道树形dp,马上推dp方程
然后发现有点复杂,好像不太对……
其实想偏了。
首先先缩点,一样颜色的缩在一起,这个不难,可以类似tarjan缩点,先求出每个点的归属,再遍历一遍树建立新的边即可。
新的图可以发现是黑白相邻的,这种情况下,从一个点开始反复操作是最优的,很容易发现最终的步数取决于最长的链,也就是直径。
所以对新图求一下树的直径即可。
#include<bits/stdc++.h>
#define rep(i, a, b) for(int i = (a); i < (b); i++)
#define _for(i, a, b) for(int i = (a); i <= (b); i++)
using namespace std;typedef long long ll;
const int N = 2e5 + 10;
int c[N], f[N], d[N], n, ans;
vector<int> g[N], G[N];int find(int x) { return x == f[x] ? x : f[x] = find(f[x]); }
void Union(int x, int y) { f[find(x)] = find(y); }void dfs1(int u, int fa)
{for(int v: g[u]){if(v == fa) continue;if(c[u] == c[v]) Union(u, v);dfs1(v, u);}
}void dfs2(int u, int fa)
{for(int v: g[u]){if(v == fa) continue;if(find(u) != find(v)) G[find(u)].push_back(find(v));dfs2(v, u);}
}void dfs3(int u, int fa)
{for(int v: G[u]){if(v == fa) continue;dfs3(v, u);ans = max(ans, d[u] + d[v] + 1);d[u] = max(d[u], d[v] + 1);}
}int main()
{scanf("%d", &n);_for(i, 1, n) scanf("%d", &c[i]), f[i] = i;_for(i, 1, n - 1){int u, v;scanf("%d%d", &u, &v);g[u].push_back(v);g[v].push_back(u);}dfs1(1, 0);dfs2(1, 0);dfs3(find(1), 0);printf("%d\n", (ans + 1) / 2);return 0;
}