結果

問題 No.1002 Twotone
ユーザー leaf_1415leaf_1415
提出日時 2020-02-28 22:41:37
言語 C++11
(gcc 11.4.0)
結果
WA  
実行時間 -
コード長 3,079 bytes
コンパイル時間 1,305 ms
コンパイル使用メモリ 79,380 KB
実行使用メモリ 46,512 KB
最終ジャッジ日時 2024-10-13 18:32:32
合計ジャッジ時間 12,580 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 4 ms
8,136 KB
testcase_01 AC 4 ms
8,084 KB
testcase_02 AC 4 ms
8,124 KB
testcase_03 AC 289 ms
16,720 KB
testcase_04 AC 389 ms
19,248 KB
testcase_05 AC 385 ms
19,268 KB
testcase_06 AC 4 ms
8,644 KB
testcase_07 WA -
testcase_08 WA -
testcase_09 WA -
testcase_10 AC 4 ms
8,128 KB
testcase_11 WA -
testcase_12 WA -
testcase_13 WA -
testcase_14 AC 4 ms
8,028 KB
testcase_15 WA -
testcase_16 WA -
testcase_17 WA -
testcase_18 AC 3 ms
7,916 KB
testcase_19 AC 429 ms
25,336 KB
testcase_20 AC 565 ms
34,756 KB
testcase_21 AC 563 ms
34,036 KB
testcase_22 AC 4 ms
8,644 KB
testcase_23 WA -
testcase_24 WA -
testcase_25 WA -
testcase_26 AC 4 ms
8,076 KB
testcase_27 AC 76 ms
17,212 KB
testcase_28 AC 130 ms
18,304 KB
testcase_29 AC 116 ms
18,916 KB
testcase_30 AC 4 ms
8,128 KB
testcase_31 WA -
testcase_32 WA -
testcase_33 WA -
testcase_34 WA -
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <iostream>
#include <algorithm>
#include <map>
#include <vector>
#include <utility>
#define llint long long
#define inf 1e18
#define mod 1000000007

using namespace std;
typedef pair<llint, llint> P;

struct edge{
	llint to, color;
	edge(){}
	edge(llint a, llint b){
		to = a, color = b;
	}
};

llint n, k;
vector<edge> G[200005];
int size[200005];
bool used[200005];
llint inv2;
llint ans = 0;

llint modpow(llint a, llint n)
{
	if(n == 0) return 1;
	if(n % 2){
		return ((a%mod) * (modpow(a, n-1)%mod)) % mod;
	}
	else{
		return modpow((a*a)%mod, n/2) % mod;
	}
}

int sizedfs(int v, int pre)
{
	int ret = 1;
	for(int i = 0; i < G[v].size(); i++){
		if(G[v][i].to == pre) continue;
		if(used[G[v][i].to]) continue;
		ret += sizedfs(G[v][i].to, v);
	}
	return size[v] = ret;
}

int centdfs(int v, int pre, int sz)
{
	for(int i = 0; i < G[v].size(); i++){
		if(G[v][i].to == pre) continue;
		if(used[G[v][i].to]) continue;
		if(size[G[v][i].to] > sz/2) return centdfs(G[v][i].to, v, sz);
	}
	return v;
}

map<P, llint> mp, tmp;
void dfs(llint v, llint p, llint c1, llint c2)
{
	P c = P(c1, c2);
	if(c.first > c.second) swap(c.first, c.second);
	tmp[c]++;
	
	for(int i = 0; i < G[v].size(); i++){
		llint u = G[v][i].to, c = G[v][i].color, C1, C2;
		if(u == p) continue;
		if(used[u]) continue;
		C1 = c1, C2 = c2;
		if(C1 != c){
			if(C2 == 0) C2 = c;
			else if(C2 != c) continue;
		}
		dfs(u, v, C1, C2);
	}
}

llint calc(map<P, llint> &mp)
{
	llint ret = 0, zsum = 0;
	for(auto it = mp.begin(); it != mp.end(); it++){
		llint x = it->second;
		if(it->first.first){
			ret += (x-1+mod)%mod * x % mod * inv2 % mod, ret %= mod;
			ret += mp[P(0LL, it->first.first)] * x % mod, ret %= mod;
			ret += mp[P(0LL, it->first.second)] * x % mod, ret %= mod;
		}
		else{
			zsum += it->second, zsum %= mod;
			ret += mod - (x-1+mod)%mod * x % mod * inv2 % mod, ret %=  mod;
		}
	}
	ret += (zsum-1+mod)%mod * zsum % mod * inv2 % mod, ret %= mod;
	return ret;
}

void solve(int v, int sz)
{
	sizedfs(v, -1);
	v = centdfs(v, -1, sz);
	
	//cout << v << endl;
	
	/* 処理  */
	mp.clear();
	for(int i = 0; i < G[v].size(); i++){
		if(used[G[v][i].to]) continue;
		tmp.clear();
		dfs(G[v][i].to, v, G[v][i].color, 0);
		ans += mod - calc(tmp), ans %= mod;
		for(auto it = tmp.begin(); it != tmp.end(); it++){
			mp[it->first] += it->second;
		}
		//cout << calc(tmp) << endl;
	}
	for(auto it = mp.begin(); it != mp.end(); it++){
		//cout << it->first.first << " " << it->first.second << " " << it->second << endl;
		if(it->first.first) ans += it->second, ans %= mod;
	}
	//cout << calc(mp) << endl;
	ans += calc(mp), ans %= mod;
	//cout << "! " << ans << endl;
	
	used[v] = true;
	for(int i = 0; i < G[v].size(); i++){
		if(used[G[v][i].to]) continue;
		solve(G[v][i].to, size[G[v][i].to]);
	}
}

int main(void)
{
	ios::sync_with_stdio(0);
	cin.tie(0);
	
	cin >> n >> k;
	llint u, v, c;
	for(int i = 1; i <= n-1; i++){
		cin >> u >> v >> c;
		G[u].push_back(edge(v, c));
		G[v].push_back(edge(u, c));
	}
	inv2 = modpow(2, mod-2);
	
	solve(1, n);
	
	cout << ans << endl;
	
	return 0;
}
0