洛谷 P1351 联合权值

题目传送门

发现这其实是一颗树,所以距离为2的点对就是父亲到儿子的儿子或同一个点的儿子们.

#include<iostream>
#include<cstdio>
#include<cstring>
#include<string>
#include<vector>
#define m 10007

using namespace std;

long long n,head[200001],tot,ans,a[200001],ans1,sum[200001],num[200001];
bool vis[200001];
struct kkk {
	int to,next;
}e1[400001];

inline long long _r() {
	long long s = 0,w = 1;
	char p = getchar();
	while(p < '0' || p > '9') {
		if(p == '-') w = -1;
		p = getchar();
	}
	while(p >= '0' && p <= '9') {
		s = s * 10 + (p - '0');
		p = getchar();
	}
	return w * s;
}

inline void add(int x,int y) {
	e1[++tot].to = y;
	e1[tot].next = head[x];
	head[x] = tot;
}

inline void dfs(int rt,int fa) {
	vis[rt] = 1;
	for(int i = head[rt];i; i = e1[i].next) {
		int u = e1[i].to;
		if(vis[u]) continue;
		ans += a[fa] * a[u],ans1 = max(ans1,a[fa] * a[u]);
		ans = ans % m;
		ans += sum[rt] * a[u],ans = ans % m;//乘法结合律优化 
		ans1 = max(ans1,num[rt] * a[u]); 
		sum[rt] += a[u];
		num[rt] = max(num[rt],a[u]);
		dfs(u,rt);
	}
}

int main() {
	n = _r();
	for(int i = 1;i < n; i++) {
		int x,y;
		x = _r();y = _r();
		add(x,y);
		add(y,x);
	}
	for(int i = 1;i <= n; i++)
		a[i] = _r();
	dfs(1,0);
	printf("%lld %lld",ans1,ans * 2 % m);
	return 0;
}