#436271 (C++11) No.1002 Twotone

• 問題: No.1002 Twotone
• ユーザー: leaf_1415
• 提出日時: 2020-02-28 22:52:41
• 言語: C++11 (gcc 11.4.0)
• 結果: WA
• 実行時間: -
• コード長: 2,654 bytes
• コンパイル時間: 1,339 ms
• コンパイル使用メモリ: 78,876 KB
• 実行使用メモリ: 46,592 KB
• 最終ジャッジ日時: 2024-10-13 18:45:21
• 合計ジャッジ時間: 11,761 ms
• ジャッジサーバーID （参考情報）: judge3 / judge2

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	4 ms; 8,064 KB
testcase_01	AC	4 ms; 8,064 KB
testcase_02	AC	4 ms; 8,064 KB
testcase_03	AC	249 ms; 16,640 KB
testcase_04	WA	-
testcase_05	AC	366 ms; 19,328 KB
testcase_06	WA	-
testcase_07	WA	-
testcase_08	WA	-
testcase_09	WA	-
testcase_10	WA	-
testcase_11	WA	-
testcase_12	WA	-
testcase_13	WA	-
testcase_14	WA	-
testcase_15	WA	-
testcase_16	WA	-
testcase_17	WA	-
testcase_18	WA	-
testcase_19	WA	-
testcase_20	WA	-
testcase_21	WA	-
testcase_22	WA	-
testcase_23	WA	-
testcase_24	WA	-
testcase_25	WA	-
testcase_26	WA	-
testcase_27	WA	-
testcase_28	WA	-
testcase_29	WA	-
testcase_30	WA	-
testcase_31	WA	-
testcase_32	WA	-
testcase_33	WA	-
testcase_34	WA	-

ソースコード

#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;
}