HJ125 最大最小路

知识点并查集

校招时部分企业笔试将禁止编程题跳出页面，为提前适应，练习时请使用在线自测，而非本地IDE。

描述

对于给定的无向无根树，第 ii 个节点上有一个权值 wiwi 。我们定义一条简单路径是好的，当且仅当：路径上的点的点权最小值小于等于 aa ，路径上的点的点权最大值大于等于 bb 。

保证给定的 a<ba<b ，你需要计算有多少条简单路径是好的。

输入描述：

第一行输入三个整数 n,a,b(1≤n≤5×105,1≤a<b≤109)n,a,b(1≤n≤5×105,1≤a<b≤109) 代表节点数、给定的上下限。

第二行输入 nn 个整数 w1,w2,...,wn(1≤wi≤109)w1,w2,...,wn(1≤wi≤109) 代表每个节点的权值。

此后 n−1n−1 行，每行输入两个整数 u,v(1≤u,v≤n,u≠v)u,v(1≤u,v≤n,u=v) 代表一条无向边连接树上 uu 和 vv 两个节点。

输出描述：

在一行上输出一个整数，代表好路径的条数。

示例1

输入：

复制代码

复制输出：

复制代码

复制说明：

复制代码

对于这个样例，如下图所示。路径 2→1→3→52→1→3→5 是好的，因为路径点权最小值 1≦a1≦a 且点权最大值 5≧b5≧b 。

复制代码

除此之外，以下路径也是好的：
∙ ∙1→3→51→3→5 ；
∙ ∙3→53→5 ；
∙ ∙4→3→54→3→5 。

cpp 复制代码

#include <iostream>
#include <vector>
#include <functional>
using namespace std;
using ll = long long;

class UnionFind {
  public:
    vector<int> parent;
    vector<int> size;

    UnionFind(int n) : parent(n + 1), size(n + 1, 1) {
        for (int i = 0; i <= n; ++i)
            parent[i] = i;
    }

    int find(int u) {
        while (parent[u] != u) {
            parent[u] = parent[parent[u]]; // 路径压缩
            u = parent[u];
        }
        return u;
    }

    void unite(int u, int v) {
        int ru = find(u), rv = find(v);
        if (ru == rv) return;
        if (size[ru] < size[rv]) swap(ru, rv);
        parent[rv] = ru;
        size[ru] += size[rv];
    }
};

ll countGoodPaths(int n, const vector<vector<int>>& tree,
                  const vector<int>& values, int a, int b) {
    // 预处理所有可能的边（防止重复处理）
    vector<pair<int, int>> edges;
    for (int u = 1; u <= n; ++u) {
        for (int v : tree[u]) {
            if (u < v) edges.emplace_back(u, v);
        }
    }

    auto calculate = [&](const function<bool(int)>& cond) -> ll {
        UnionFind uf(n);
        vector<bool> valid(n + 1);
        for (int i = 1; i <= n; ++i)
            valid[i] = cond(i);

        // 合并符合条件的边
        for (auto& [u, v] : edges) {
            if (valid[u] && valid[v])
                uf.unite(u, v);
        }

        // 统计连通块大小
        vector<int> cnt(n + 1, 0);
        for (int i = 1; i <= n; ++i) {
            if (valid[i])
                cnt[uf.find(i)]++;
        }

        // 计算总路径数
        ll sum = 0;
        for (int c : cnt) {
            if (c >= 2)
                sum += (ll)c * (c - 1) / 2;
        }
        return sum;
    };

    const ll total = (ll)n * (n - 1) / 2;
    const ll max_less_b = calculate([&](int x) {
        return values[x] < b;
    });
    const ll min_greater_a = calculate([&](int x) {
        return values[x] > a;
    });
    const ll both = calculate([&](int x) {
        return values[x] > a && values[x] < b;
    });

    return total - max_less_b - min_greater_a + both;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n, a, b;
    cin >> n >> a >> b;

    vector<int> values(n + 1); // values[1..n]
    for (int i = 1; i <= n; ++i)
        cin >> values[i];

    vector<vector<int>> tree(n + 1);
    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        cin >> u >> v;
        tree[u].push_back(v);
        tree[v].push_back(u);
    }

    cout << countGoodPaths(n, tree, values, a, b) << endl;

    return 0;
}