题解 | 最大最小路

#include<iostream>
#include<vector>
#include<algorithm>
#include<map>
#include<string>
#include<cmath>
#include<queue>
#include<sstream>
#include<iomanip>

using namespace std;

long long findrootparent(long long u, vector<long long>& parent)
{
	while (parent[u] != u)
	{
		u = parent[u];
	}
	return u;
}

long long calroutes(vector<pair<long long, long long>>& lines, vector<bool>& valid)
{
	long long n = valid.size();
	vector<long long> parent(n);
	for (long long i = 0; i < n; i++)
	{
		parent[i] = i;
	}
	vector<long long> leafnum(n, 1);
	
	//cout << "n - 1 = " << n - 1 << " lines.size() = " << lines.size() << endl;
	for (long long i = 0; i < n - 1; i++)
	{
		long long u = lines[i].first;
		long long v = lines[i].second;
		
		if (!valid[u] || !valid[v])
			continue;
	
		long long rootu = findrootparent(u, parent);
		long long rootv = findrootparent(v, parent);
		if (leafnum[rootu] < leafnum[rootv])
			swap(rootu, rootv);

		parent[rootv] = rootu;
		leafnum[rootu] = leafnum[rootu] + leafnum[rootv];
	}
	
	map<long long, long long> recordsroots;
	for (long long i = 0; i < n; i++)
	{
		long long root = findrootparent(i, parent);
		recordsroots[root] = 1;
	}

	long long countroutes = 0;
	for (auto it : recordsroots)
	{
		long long numpoints = leafnum[it.first];
		countroutes = countroutes + numpoints * (numpoints - 1) / 2;
	}

	return countroutes;
}

int main()
{
	long long n, a, b;
	cin >> n >> a >> b;
	vector<long long> ws(n);
	vector<bool> validmorethana(n, false);
	vector<bool> validlessthanb(n, false);
	vector<bool> validbetweenab(n, false);
	for (long long i = 0; i < n; i++)
	{
		cin >> ws[i];
		if (ws[i] > a)
			validmorethana[i] = true;
		if (ws[i] < b)
			validlessthanb[i] = true;
		if (ws[i] > a && ws[i] < b)
			validbetweenab[i] = true;
	}
	vector<pair<long long, long long>> lines(n - 1);
	for (long long i = 0; i < n - 1; i++)
	{
		cin >> lines[i].first >> lines[i].second;
		lines[i].first--;
		lines[i].second--;
	}

	long long routesmorethana = calroutes(lines, validmorethana);
	long long routeslessthanb = calroutes(lines, validlessthanb);
	long long routesbetweenab = calroutes(lines, validbetweenab);

	long long routesvalid = n * (n - 1) / 2 - (routesmorethana + routeslessthanb - routesbetweenab);
	cout << routesvalid << endl;

	return 0;
}