题解 | 最大最小路

#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;
}