E题没有数据找不到WA的点在那 求改 只通过了30%的测试点

#include <bits/stdc++.h>

using namespace std;
#define int long long
#define BUG cout << "HERE IS BUG !!!\n"
#define endl '\n'
const int N = 1e6 + 1;

#define inf 1e18
int n;
int type[N];
int w[N];
int dp1[N], dp2[N];//dp1是最大路径点权值 dp2是次大路径点权值
int ans[N];
vector<int> g[N];
void dfs(int u, int fa){
	dp1[u] = w[u], dp2[u] = w[u];
	for (auto v : g[u]){
		if (v == fa) continue;
		dfs(v, u);
		if (dp1[v] + w[u] > dp1[u]) dp2[u] = dp1[u], dp1[u] = dp1[v] + w[u];
		else if (dp1[v] + w[u] > dp2[u]) dp2[u] = dp1[v] + w[u];
	}
}


void dfs1(int u, int fa){
	
	for (auto v : g[u]){
		if (v == fa) continue;
		int dp1_u = dp1[u], dp2_u = dp2[u], dp1_v = dp1[v], dp2_v = dp2[v];
		//树根由u转为v后树根v的最大路径点权和次大路径点权
		if (dp1_v + w[u] != dp1_u && dp1_v + w[u] != dp2_u){
			dp1[v] = max(dp1_v, dp1_u + w[v]);
			dp2[v] = min(dp1_v, dp1_u + w[v]);
		}else if (dp1_v + w[u] == dp1_u){
			dp1[v] = max(dp1_v, dp2_u + w[v]);
			dp2[v] = max(dp2_v, min(dp1_v, dp2_u + w[v]));
		}else if (dp1_v + w[u] == dp2_u){
			dp1[v] = max(dp1_v, dp1_u + w[v]);
			dp2[v] = max(dp2_v, min(dp1_v, dp1_u + w[v]));
		}
		
		//更新答案
		int o = max(dp1[v] + dp2[v] - w[v], max(dp1[v], dp2[v]));
		ans[type[v]] = max(ans[type[v]], o);
		dfs1(v, u);
	}
}


void solve(){
	cin >> n;
	for (int i = 1; i <= n; i ++) cin >> type[i], g[i].clear();
	for (int i = 1; i <= n; i ++) cin >> w[i], dp1[i] = dp2[i] = 0, ans[i] = -inf;
	
	for (int i = 1; i < n; i ++){
		int u, v; cin >> u >> v;
		g[u].push_back(v); 
		g[v].push_back(u);
	}
	
	dfs(1, - 1);
	
	ans[type[1]] = max(dp1[1] + dp2[1] - w[1], max(dp1[1], dp2[1]));

	dfs1(1, -1);			
	
	for (int i = 1; i <= n; i ++){
		if (ans[i] == -inf) cout << -1 << ' ';
		else cout << ans[i] << ' ';
	}
	
	cout << "\n";
}


signed main(){
	ios::sync_with_stdio(false), cin.tie(0), cout.tie(0);
	int t = 1;
	cin >> t;
	while (t --) solve();
	
}

全部评论

相关推荐

评论
点赞
收藏
分享

创作者周榜

更多
牛客网
牛客企业服务