H. Crystalfly（2021ICPC-南京）（树状dp）

The 2021 ICPC Asia Nanjing Regional Contest
参考题解
比赛题解

Solution

Code

const int N =  2e5 + 5;
ll sum[N], f[N], a[N], t[N];
vector<int> v[N];

void dfs(int x, int fa)
{
	multiset<ll>s;//set中根节点是集合最小值
	ll maxn = 0;
	for (int y : v[x])
	{
		if (y == fa) continue;
		dfs(y, x);
		
		sum[x] += f[y];
		maxn = max(maxn, a[y]);
		if (t[y] == 3) s.insert(a[y]);
	}
	f[x] = sum[x] + maxn;

	s.insert(-1e18);//防止对空集的访问
	for (int y : v[x])
	{
		if (y == fa) continue;
		if (t[y] == 3) s.erase(s.find(a[y]));
		f[x] = max(f[x], sum[x] - f[y] + sum[y] + *s.rbegin() + a[y]);
		//注意是rbegin，end会报错（end指向最后一个元素的下一个位置）
		if(t[y] == 3) s.insert(a[y]);
	}
}

int main()
{
	IOS;
	int T; cin >> T;
	while (T--)
	{
		int n; cin >> n;			
		for (int i = 1; i <= n; i++)
		{
			sum[i] = f[i] = 0;
			v[i].clear();
			cin >> a[i];
		}
		for (int i = 1; i <= n; i++)
			cin >> t[i];

		for (int i = 1, x, y; i < n; i++)
		{
			cin >> x >> y;
			v[x].push_back(y);
			v[y].push_back(x);
		}

		dfs(1, -1);
		cout << a[1] + f[1] << endl;
	}

	return 0;
}