結果

問題 No.1002 Twotone
コンテスト
ユーザー leaf_1415
提出日時 2020-02-28 22:53:49
言語 C++11（廃止可能性あり）
(gcc 13.3.0)
結果
AC  
実行時間 1,120 ms / 5,000 ms
コード長 2,654 bytes
コンパイル時間 706 ms
コンパイル使用メモリ 79,904 KB
実行使用メモリ 46,464 KB
最終ジャッジ日時 2024-10-13 18:45:57
合計ジャッジ時間 10,039 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 33
権限があれば一括ダウンロードができます

ソースコード

diff # 

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

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 ans = 0;

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 get(llint x)
{
	return x*(x-1)/2;
}

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 += get(x);
			ret += mp[P(0LL, it->first.first)] * x;
			ret += mp[P(0LL, it->first.second)] * x;
		}
		else{
			zsum += x;
			ret -= get(x);
		}
	}
	ret += get(zsum);
	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 -= calc(tmp);
		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;
	}
	//cout << calc(mp) << endl;
	ans += calc(mp);
	//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));
	}
	
	solve(1, n);
	
	cout << ans << endl;
	
	return 0;
}
0