11111

#include 

#define int long long

constexpr int N = 2e5 + 10;

std::vector adj[N];

int n;
int a[N];
int lst[N];
int sz[N], dep[N], F[N][33];
int dp[N];

void dfs(int u, int fa) {
	sz[u] = 1;
	dep[u] = dep[fa] + 1;
	F[u][0] = fa;
	for (int i = 1; i <= 30; i++) {
		F[u][i] = F[F[u][i - 1]][i - 1];
	}
	for (auto v : adj[u]) {
		if (v == fa) continue;
		dfs(v, u);
		sz[u] += sz[v];
	}
}
void dfs2(int u, int fa) {
	if (a[u] == 1) dp[u] = 1;
	for (auto v : adj[u]) {
		if (v == fa) continue;
		dfs2(v, u);
		if (a[v] == 1) {
			dp[u] += dp[v];
		}
	}
}
void dfs3(int u, int fa, int last) {
	lst[u] = last;
	for (auto v : adj[u]) {
		if (v == fa) continue;
		int clst = 0;
		if (a[v]) {
			clst = last;
		}else {
			clst = v;
		}
		dfs3(v, u, clst); 
	}
}
int lca(int u, int v) {
	if (dep[u] < dep[v]) {
		std::swap(u, v);
	}
	for (int j = 30; j >= 0; j --) {
		if (dep[F[u][j]] >= dep[v]) {
			u = F[u][j];
		} 
	}
	if (u == v) return u;
	for (int j = 30; j >= 0; j --) {
		if (F[u][j] != F[v][j]) {
			u = F[u][j];
			v = F[v][j];
		}
	}
	return F[u][0];
}
int dis(int u, int v) {
	return dep[u] + dep[v] - 2 * dep[lca(u, v)];
}

void solve() {
	std::cin >> n;
	for (int i = 1; i <= n; i ++) {
		std::cin >> a[i];
	}
	for (int i = 1; i <= n - 1; i ++) {
		int u, v;
		std::cin >> u >> v;
		adj[u].push_back(v);
		adj[v].push_back(u);
	}
	dfs(1, 0);
	dfs3(1, 0, 1);
	for (int i = 1; i <= n; i ++) {
		lst[i] = dis(i, lst[i]);
	}
	dfs2(1, 0);
	int cur = 0;
	for (int u = 1; u <= n; u ++) {
		if (a[u] == 1) {
			cur += lst[u];
		}
	}
	int ans = cur;
	for (int u = 1; u <= n; u ++) {
		if (!a[u]) continue;
		int res = cur;
		res -= dp[u] * lst[u];
		ans = std::min(ans, res);
	}
	std::cout << ans << "\n";
}
signed main() {
	std::ios::sync_with_stdio(false);
	std::cin.tie(nullptr);
	
	int t = 1;
	while(t --) {
		solve();
	}
	return 0;
} 

你可能感兴趣的:(深度优先,算法)