华为机试—最大最小路
题目
对于给定的无向无根树，第 i 个节点上有一个权值 wi 。我们定义一条简单路径是好的，当且仅当：路径上的点的点权最小值小于等于 a ，路径上的点的点权最大值大于等于 b 。
保证给定的 a<b，你需要计算有多少条简单路径是好的。
示例
输入：
5 2 3 5 4 3 3 1 1 2 1 3 3 4 3 5第一行输入三个整数 n,a,b(1≤n≤5×
,1≤a<b≤
) 代表节点数、给定的上下限。
第二行输入 n 个整数 w1,w2,…,wn(1≤wi≤
此后 n−1 行，每行输入两个整数 u,v(1≤u,v≤n,uv) 代表一条无向边连接树上 u 和 v 两个节点。
输出：4
说明：对于这个样例，如下图所示。路径 2→1→3→5 是好的，因为路径点权最小值 1≦a 且点权最大值 5≧b。
除此之外，以下路径也是好的： ∙1→3→5； ∙3→5； ∙4→3→5。
分析
并查集+容斥原理
算法思路
统计总路径数：
根据组合数学得到所有可能的简单路径数。n 个节点的树中总路径数（包含单节点路径）可由组合公式计算：
分情况统计不好路径：
如果一条路径不好，则可能满足：
-
全部节点权值均大于 a ——此时路径的最小值大于 a；
-
全部节点权值均小于 b ——此时路径的最大值小于 b。
利用并查集对满足特定条件（全部节点权值大于 a，或全部节点权值小于 b，以及同时满足这两者的情况）构成的子图进行连通分量划分，进而统计每个连通分量内部所有路径数量。
利用容斥原理：
计算好路径数，确保重复扣除部分得到校正，从而求解出最终的答案。
时间复杂度：O()
空间复杂度：O()
#include <iostream>
#include <vector>
#include <utility>
using namespace std;
using ll = long long;
struct DSU {
vector<int> parent, size;
DSU(int n): parent(n+1), size(n+1, 1) {
for (int i = 0; i <= n; ++i) parent[i] = i;
}
int find(int x) {
return parent[x] == x ? x : parent[x] = find(parent[x]);
}
void unite(int a, int b) {
a = find(a), b = find(b);
if(a == b) return;
if(size[a] < size[b]) swap(a, b);
parent[b] = a;
size[a] += size[b];
}
};
int main(){
int n, a, b;
cin >> n >> a >> b;
vector<int> w(n+1);
for (int i = 1; i <= n; ++i) cin >> w[i];
vector<pair<int,int>> edges(n-1);
for (int i = 0; i < n-1; ++i)
cin >> edges[i].first >> edges[i].second;
// 总路径数（单节点路径也算）
ll total = (ll)n * (n+1) / 2;
auto countComponentPaths = [&](auto &dsu, const vector<bool> &valid) -> ll {
vector<int> comp(n+1, 0);
for (int i = 1; i <= n; ++i)
if(valid[i])
comp[dsu.find(i)]++;
ll sum = 0;
for (int cnt : comp)
if(cnt) sum += (ll)cnt * (cnt+1) / 2;
return sum;
};
// DSU1：构造仅包含权值 > a 的子图
DSU dsu1(n);
vector<bool> valid1(n+1, false);
for (int i = 1; i <= n; ++i)
valid1[i] = (w[i] > a);
for (auto &e : edges) {
if(valid1[e.first] && valid1[e.second])
dsu1.unite(e.first, e.second);
}
// DSU2：构造仅包含权值 < b 的子图
DSU dsu2(n);
vector<bool> valid2(n+1, false);
for (int i = 1; i <= n; ++i)
valid2[i] = (w[i] < b);
for (auto &e : edges) {
if(valid2[e.first] && valid2[e.second])
dsu2.unite(e.first, e.second);
}
// DSU3：构造仅包含 a < 权值 < b 的子图
DSU dsu3(n);
vector<bool> valid3(n+1, false);
for (int i = 1; i <= n; ++i)
valid3[i] = (w[i] > a && w[i] < b);
for (auto &e : edges) {
if(valid3[e.first] && valid3[e.second])
dsu3.unite(e.first, e.second);
}
ll cnt1 = countComponentPaths(dsu1, valid1); // 所有节点均 > a 的路径
ll cnt2 = countComponentPaths(dsu2, valid2); // 所有节点均 < b 的路径
ll cnt3 = countComponentPaths(dsu3, valid3); // 同时均在 (a,b) 内 的路径
// 由于好路径要求最小值<= a 且最大值>= b，因此
// 不好路径：所有节点均 > a 或所有节点均 < b
// 但区间在 (a, b) 的部分被重复扣除，需加回来
ll ans = total - cnt1 - cnt2 + cnt3;
cout << ans << "\n";
return 0;
}