yukicoder : No.1333 Squared Sum

これは典型かな。
https://yukicoder.me/problems/no/1333

問題

木を成す無向グラフが与えられる。
各辺には長さが設定されている。
頂点対における最短距離の二乗和を求めよ。

解法

あるパスの距離がAとする。そこの端に長さBのパスを連結すると、二乗和は(A+B)^2=A^2+B^2+2ABとなる。
そこで、各点においてSubTreeにおける

• パス数＝頂点数
• 長さの1乗和
• 長さの2乗和

を順次DPで更新していこう。
全方位DPも不要。

int N;
vector<pair<int,int>> E[202020];
const ll mo=1000000007;
ll ret;

vector<ll> dfs(int cur,int pre) {
	vector<ll> p={1,0,0};
	int pw=0;
	FORR2(e,w,E[cur]) {
		if(e==pre) {
			pw=w;
		}
		else {
			vector<ll> q=dfs(e,cur);
			(ret+=p[0]*q[2]+p[2]*q[0]+2*p[1]*q[1])%=mo;
			(p[0]+=q[0])%=mo;
			(p[1]+=q[1])%=mo;
			(p[2]+=q[2])%=mo;
			
		}
	}
	(p[2]+=p[0]*pw%mo*pw+2*p[1]*pw)%=mo;
	(p[1]+=p[0]*pw)%=mo;
	return p;
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N-1) {
		cin>>x>>y>>r;
		E[x-1].push_back({y-1,r});
		E[y-1].push_back({x-1,r});
	}
	dfs(0,0);
	cout<<ret<<endl;
}

まとめ

これは★3でもいい気はする。