O(n) 做法

进行一个树形 DP

记录向下路径是否包含共四种情况，DP 合并过程中统计答案

#include <array>
#include <iostream>
#include <vector>
using namespace std;

using ll = long long;

int n;
int a;
int b;
vector<int> w;
vector<vector<int>> graph;

int F(int x) {
    if (w[x] <= a) {
        return 1;
    }
    if (w[x] >= b) {
        return 2;
    }
    return 0;
}

ll res = 0;

array<int, 4> Dp(int x, int fa) {
    array<int, 4> ret = {};
    int v = F(x);
    ret[v]++;
    // cout << x << ',' << v << endl;
    for (const auto& y : graph[x]) {
        if (y != fa) {
            auto&& solo = Dp(y, x);
            res += (ll)ret[0] * solo[3] + (ll)ret[1] * (solo[2] + solo[3]) + (ll)ret[2] * (solo[1] + solo[3]) + (ll)ret[3] * (solo[0] + solo[1] + solo[2] + solo[3]);
            for (int i = 0; i < 4; i++) {
                ret[i | v] += solo[i];
            }
        }
    }
    // cout << x << ',' << ret[0] << ',' << ret[1] << ',' << ret[2] << ',' << ret[3] << endl;
    return ret;
}

void Solve() {
    cin >> n >> a >> b;
    w.resize(n + 1);
    graph.resize(n + 1);
    for (int i = 1; i <= n; i++) {
        cin >> w[i];
    }
    for (int i = 1; i < n; i++) {
        int u, v;
        cin >> u >> v;
        graph[u].push_back(v);
        graph[v].push_back(u);
    }
    Dp(1, 0);
    cout << res;
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    Solve();
}
// 64 位输出请用 printf("%lld")