题目
对于给定的无向无根树,第 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;
}