結果

問題 No.1002 Twotone
ユーザー leaf_1415leaf_1415
提出日時 2020-02-28 22:53:49
言語 C++11
(gcc 11.4.0)
結果
AC  
実行時間 1,114 ms / 5,000 ms
コード長 2,654 bytes
コンパイル時間 1,008 ms
コンパイル使用メモリ 78,740 KB
実行使用メモリ 46,336 KB
最終ジャッジ日時 2024-04-21 20:04:48
合計ジャッジ時間 11,037 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge2
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 5 ms
8,064 KB
testcase_01 AC 5 ms
8,064 KB
testcase_02 AC 5 ms
7,936 KB
testcase_03 AC 232 ms
16,640 KB
testcase_04 AC 326 ms
19,328 KB
testcase_05 AC 338 ms
19,328 KB
testcase_06 AC 5 ms
8,064 KB
testcase_07 AC 132 ms
14,848 KB
testcase_08 AC 241 ms
19,328 KB
testcase_09 AC 242 ms
19,200 KB
testcase_10 AC 5 ms
7,936 KB
testcase_11 AC 258 ms
17,920 KB
testcase_12 AC 365 ms
20,736 KB
testcase_13 AC 385 ms
20,736 KB
testcase_14 AC 4 ms
8,064 KB
testcase_15 AC 170 ms
15,232 KB
testcase_16 AC 375 ms
20,480 KB
testcase_17 AC 369 ms
20,480 KB
testcase_18 AC 5 ms
8,064 KB
testcase_19 AC 354 ms
25,472 KB
testcase_20 AC 454 ms
34,688 KB
testcase_21 AC 464 ms
33,920 KB
testcase_22 AC 5 ms
8,064 KB
testcase_23 AC 322 ms
21,888 KB
testcase_24 AC 618 ms
31,360 KB
testcase_25 AC 590 ms
30,080 KB
testcase_26 AC 5 ms
8,064 KB
testcase_27 AC 65 ms
15,932 KB
testcase_28 AC 110 ms
18,236 KB
testcase_29 AC 96 ms
18,364 KB
testcase_30 AC 4 ms
8,064 KB
testcase_31 AC 96 ms
18,104 KB
testcase_32 AC 112 ms
18,236 KB
testcase_33 AC 94 ms
18,104 KB
testcase_34 AC 1,114 ms
46,336 KB
権限があれば一括ダウンロードができます

ソースコード

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