結果

問題 No.1333 Squared Sum
ユーザー 👑 potato167potato167
提出日時 2022-02-25 17:05:53
言語 C++17
(gcc 13.3.0 + boost 1.87.0)
結果
RE  
実行時間 -
コード長 2,792 bytes
コンパイル時間 2,401 ms
コンパイル使用メモリ 206,076 KB
最終ジャッジ日時 2025-01-28 01:46:00
ジャッジサーバーID
（参考情報） 		judge2 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other AC * 1 RE * 43
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <bits/stdc++.h>
using namespace std;
using std::cout;
using std::cin;
using std::endl;
using ll=long long;
ll mod=1e9+7;
#define rep(i,a) for (ll i=0;i<a;i++)


namespace po167{
struct UFtree
{
	std::vector<int> wei;
	int component;
	UFtree(int n):par(n),wei(n),component(n){
		for(int i=0;i<n;i++){
			wei[i]=1,par[i]=i;
		}
	}
	void intialize(){
		for(int i=0;i<(int) par.size();i++){
			wei[i]=1,par[i]=i;
		}
		component=(int)par.size();
	}
	//根っこを返す
	int root(int a){
		if(a==par[a]) return a;
		return par[a]=root(par[a]);
	}
	//trueなら1,falseなら0
	int same(int a,int b){
		if(root(a)==root(b)) return 1;
		else return 0;
	}
	//a,bが違う根っこの元なら結合する,結合したらtrueを返す
	bool unite(int a,int b){
		a=root(a),b=root(b);
		if(a==b) return false;
		if(wei[a]<wei[b]) std::swap(a,b);
		par[b]=a;
		wei[a]+=wei[b];
		component--;
		return true;
	}
	private:
	std::vector<int> par;
};
}
using po167::UFtree;

struct edge{
	int to;
	ll len;
};

const int MAX_N=2e5;
const int MAX_C=1e9;

//おちこんだりもしたけれど、私はげんきです。
int main() {
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	
	ll N;
	cin>>N;
	assert(2<=N&&N<=MAX_N);
	UFtree tree(N);
	vector<vector<edge>> G(N);

	rep(i,N-1){
		int a,b;ll c;
		cin>>a>>b>>c;
		assert(1<=c&&c<=MAX_C);
		assert(1<=a&&a<b&&b<=N);
		a--,b--;
		assert(tree.unite(a,b));
		G[a].push_back({b,c});
		G[b].push_back({a,c});
	}

	ll ans=0;
	vector<ll> subtree_size(N,1);
	vector<edge> parent(N,{-1,-1});
	vector<ll> order={0};
	vector<ll> lower_dist_sum(N);
	vector<ll> upper_dist_sum(N);

	parent[0]={-2,-2};
	rep(i,N){
		int var=order[i];
		for(auto e:G[var]){
			if(parent[e.to].to==-1){
				parent[e.to]={var,e.len};
				order.push_back(e.to);
			}
		}
	}
	for(int i=N-1;i>=0;i--){
		int var=order[i];
		for(auto e:G[var]){
			if(e.to!=parent[var].to){
				subtree_size[var]+=subtree_size[e.to];
				lower_dist_sum[var]+=(subtree_size[e.to]*e.len)%mod;
				lower_dist_sum[var]+=lower_dist_sum[e.to];
				lower_dist_sum[var]%=mod;
			}
		}
	}

	rep(i,N){
		int var=order[i];
		for(auto e:G[var]){
			if(e.to==parent[var].to) continue;
			ll tmp_lower=(lower_dist_sum[e.to]+(subtree_size[e.to]*e.len)%mod)%mod;
			upper_dist_sum[e.to]=(upper_dist_sum[var]+lower_dist_sum[var]-tmp_lower)%mod;
			ll ans_sub=(upper_dist_sum[e.to]*subtree_size[e.to]);
			ans_sub+=(tmp_lower*(N-subtree_size[e.to]))%mod;
			ans_sub=((ans_sub%mod)*e.len)%mod;
			//cout<<ans_sub<<endl<<endl;
			ans+=ans_sub;
			ans%=mod;
			upper_dist_sum[e.to]+=(e.len*(N-subtree_size[e.to]))%mod;
			upper_dist_sum[e.to]%=mod;
		}
	}

	/*rep(i,N){
		cout<<upper_dist_sum[i]<<" "<<lower_dist_sum[i]<<" "<<subtree_size[i]<<endl;
		cout<<parent[i].to<<" "<<parent[i].len<<endl;
	}*/
	cout<<ans<<endl;
}
0