华为机试—最大最小路

题目

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