发现一个超简单的解法

#include <bits/stdc++.h>
#define int long long
using namespace std;

const int N = 1e5 + 5;

int n, ans, root;
int a[N];
vector <int> g[N];

void DFS(int u, int fa) {
	for(auto &v : g[u]) {
		if(v == fa) continue;
		DFS(v, u);
		if(a[v] == 1 && a[u] == 0) a[u] = 1, ans++;
	}
}

signed main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr), cout.tie(nullptr);
	cin >> n;
	for(int i = 1; i <= n; i++) {
		cin >> a[i];
		if(a[i]) root = i;
	}
	for(int i = 1; i < n; i++) {
		int u, v;
		cin >> u >> v;
		g[u].emplace_back(v), g[v].emplace_back(u);
	}
	DFS(root, 0);
	cout << ans << endl;
	return 0;
}