O(n) 做法
最大最小路
https://www.nowcoder.com/practice/bce59fc8643b47359e9521c7e590b239
进行一个树形 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")
