容斥原理 + 并查集维护
#include <bits/stdc++.h>
#define x first
#define y second
#define all(x) x.begin(), x.end()
#define vec1(T, name, n, val) vector<T> name(n, val)
#define vec2(T, name, n, m, val) vector<vector<T>> name(n, vector<T>(m, val))
#define vec3(T, name, n, m, k, val) vector<vector<vector<T>>> name(n, vector<vector<T>>(m, vector<T>(k, val)))
#define vec4(T, name, n, m, k, p, val) vector<vector<vector<vector<T>>>> name((n), vector<vector<vector<T>>>((m), vector<vector<T>>((k), vector<T>((p), (val)))))
using namespace std;
using i128 = __int128;
using u128 = unsigned __int128;
using LL = long long;
using LD = long double;
using ULL = unsigned long long;
using PII = pair<int, int>;
using PLL = pair<LL, LL>;
using PLD = pair<LD, LD>;
const int N = 1e5 + 10, MOD = 998244353;
const int INF = 1e9;
const LL LL_INF = 2e18;
const LD EPS = 1e-8;
const int dx4[] = {-1, 0, 1, 0}, dy4[] = {0, 1, 0, -1};
const int dx8[] = {-1, -1, -1, 0, 0, 1, 1, 1}, dy8[] = {-1, 0, 1, -1, 1, -1, 0, 1};
istream& operator>>(istream& is, i128& val) {
string str;
is >> str;
val = 0;
bool flag = false;
if (str[0] == '-') flag = true, str = str.substr(1);
for (char& c : str) val = val * 10 + c - '0';
if (flag) val = -val;
return is;
}
ostream& operator<<(ostream& os, i128 val) {
if (val < 0) os << "-", val = -val;
if (val > 9) os << val / 10;
os << static_cast<char>(val % 10 + '0');
return os;
}
bool cmp(LD a, LD b) {
if (fabs(a - b) < EPS) return 1;
return 0;
}
LL qpow(LL a, LL b) {
LL ans = 1;
a %= MOD;
while (b) {
if (b & 1) ans = ans * a % MOD;
a = a * a % MOD;
b >>= 1;
}
return ans;
}
struct DSU {
int n;
vector<int> p;
vector<int> sz;
DSU(int _n) : n(_n), p(_n), sz(_n) {
for (int i = 0; i < _n; ++i) {
p[i] = i;
}
for (int i = 0; i < _n; ++i) sz[i] = 1;
};
int find(int x) {
if (p[x] != x) return p[x] = find(p[x]);
return p[x];
}
void merge(int x, int y) {
int fa1 = find(x), fa2 = find(y);
if (fa1 == fa2) return;
p[fa2] = fa1;
sz[fa1] += sz[fa2];
}
};
void solve() {
int n, a, b;
cin >> n >> a >> b;
vector<int> w(n + 1, 0);
for (int i = 1; i <= n; ++i) cin >> w[i];
vector<PII> ed(n - 1);
for (int i = 0; i < n - 1; ++i) cin >> ed[i].x >> ed[i].y;
auto calc = [&](LL l, LL r, vector<PII>& vec) {
if (l > r) return 0ll;
DSU d(n + 1);
for (auto &[x, y] : vec) {
if (w[x] >= l && w[x] <= r && w[y] >= l && w[y] <= r) d.merge(x, y);
}
LL res = 0;
for (int i = 1; i <= n; ++i) {
if (w[i] >= l && w[i] <= r && d.find(i) == i) {
LL cnt = d.sz[i];
res += (cnt + 1) * cnt / 2;
}
}
return res;
};
LL ans = 1ll * (n + 1) * n / 2;
ans -= calc(a + 1, LL_INF, ed);
ans -= calc(-LL_INF, b - 1, ed);
ans += calc(a + 1, b - 1, ed);
cout << ans << '\n';
/**/ #ifdef LOCAL
cout
<< flush;
/**/ #endif
}
int main() {
ios::sync_with_stdio(false);
cin.tie(0), cout.tie(0);
int T = 1;
while (T--) solve();
cout << fixed << setprecision(15);
return 0;
}

京公网安备 11010502036488号