#437209 (C++14) No.1002 Twotone

• 問題: No.1002 Twotone
• ユーザー: QCFium
• 提出日時: 2020-03-02 22:28:08
• 言語: C++14 (gcc 12.3.0 + boost 1.83.0)
• 結果: TLE
• 実行時間: -
• コード長: 3,165 bytes
• コンパイル時間: 3,015 ms
• コンパイル使用メモリ: 198,148 KB
• 実行使用メモリ: 203,524 KB
• 最終ジャッジ日時: 2024-04-21 22:53:39
• 合計ジャッジ時間: 17,605 ms
• ジャッジサーバーID （参考情報）: judge5 / judge2

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	2 ms; 5,248 KB
testcase_01	AC	2 ms; 5,376 KB
testcase_02	AC	2 ms; 5,376 KB
testcase_03	AC	436 ms; 14,848 KB
testcase_04	AC	686 ms; 18,176 KB
testcase_05	AC	627 ms; 18,176 KB
testcase_06	AC	2 ms; 5,376 KB
testcase_07	AC	248 ms; 12,288 KB
testcase_08	AC	467 ms; 17,920 KB
testcase_09	AC	485 ms; 17,920 KB
testcase_10	AC	2 ms; 5,376 KB
testcase_11	AC	461 ms; 16,128 KB
testcase_12	AC	718 ms; 19,968 KB
testcase_13	AC	806 ms; 20,224 KB
testcase_14	AC	2 ms; 5,376 KB
testcase_15	AC	321 ms; 12,416 KB
testcase_16	AC	810 ms; 18,816 KB
testcase_17	AC	779 ms; 19,072 KB
testcase_18	AC	3 ms; 5,376 KB
testcase_19	TLE	-
testcase_20	--	-
testcase_21	--	-
testcase_22	--	-
testcase_23	--	-
testcase_24	--	-
testcase_25	--	-
testcase_26	--	-
testcase_27	--	-
testcase_28	--	-
testcase_29	--	-
testcase_30	--	-
testcase_31	--	-
testcase_32	--	-
testcase_33	--	-
testcase_34	--	-

ソースコード

#include <bits/stdc++.h>

int ri() {
	int n;
	scanf("%d", &n);
	return n;
}

std::vector<std::vector<std::pair<int, int> > > hen;
std::vector<bool> alive;
std::vector<int> size;
void dfs(int i, int prev) {
	size[i] = 1;
	for (auto j : hen[i]) if (j.first != prev && alive[j.first]) {
		dfs(j.first, i);
		size[i] += size[j.first];
	}
}
int get_centroid(int i) {
	dfs(i, -1);
	while (1) {
		int max = 0;
		int maxindex = -1;
		for (auto j : hen[i]) if (alive[j.first] && size[j.first] < size[i] && max < size[j.first])
			max = size[j.first], maxindex = j.first;
		if (max > size[i] / 2) i = maxindex;
		else break;
	}
	return i;
}
int64_t res = 0;
std::vector<std::pair<int, int> > path_color;
void centroid_rec(int i) {
	i = get_centroid(i);
	struct Info {
		std::map<std::pair<int, int>, int> two_match;
		std::map<int, int> two_side;
		std::map<int, int> one;
		int one_num = 0;
	};
	std::vector<Info> infos;
	for (auto j : hen[i]) {
		if (!alive[j.first]) continue;
		infos.push_back({});
		Info &info = infos.back();
		std::queue<std::pair<int, int> > que;
		path_color[j.first] = {j.second, -1};
		que.push({j.first, i});
		while (que.size()) {
			auto k = que.front();
			auto &cur_color = path_color[k.first];
			que.pop();
			if (cur_color.second != -1) {
				if (cur_color.second < cur_color.first) std::swap(cur_color.first, cur_color.second);
				info.two_match[cur_color]++;
				info.two_side[cur_color.first]++;
				info.two_side[cur_color.second]++;
			} else {
				info.one_num++;
				info.one[cur_color.first]++;
			}
			for (auto l : hen[k.first]) if (l.first != k.second && alive[l.first]) {
				path_color[l.first] = cur_color;
				if (path_color[l.first].first != l.second) {
					if (path_color[l.first].second == -1) path_color[l.first].second = l.second;
					else if (path_color[l.first].second != l.second) continue;
				}
				que.push({l.first, k.first});
			}
		}
	}
	Info all;
	for (auto &j : infos) {
		for (auto &k : j.two_match) all.two_match[k.first] += k.second;
		for (auto &k : j.two_side) all.two_side[k.first] += k.second;
		for (auto &k : j.one) all.one[k.first] += k.second;
		all.one_num += j.one_num;
	}
	auto calc = [&] (Info &target, Info &main) {
		int64_t res = 0;
		for (auto &j : main.two_match) {
			res += (int64_t) j.second * target.two_match[j.first];
			res += (int64_t) j.second * (target.one[j.first.first] + target.one[j.first.second]);
		}
		for (auto &j : main.one) {
			res += (int64_t) j.second * target.two_side[j.first];
			res += (int64_t) j.second * (target.one_num - target.one[j.first]);
		}
		return res;
	};
	for (auto &j : infos) res += calc(all, j) - calc(j, j);
	for (auto &j : all.two_match) res += j.second * 2;
	alive[i] = false;
	for (auto &j : hen[i]) if (alive[j.first]) centroid_rec(j.first);
}

void run() {
	int n = hen.size();
	alive.resize(n,	true);
	size.resize(n);
	path_color.resize(n);
	centroid_rec(0);
	res /= 2;
}

int main() {
	int n = ri();
	int k = ri();
	hen.resize(n);
	for (int i = 1; i < n; i++) {
		int x = ri() - 1;
		int y = ri() - 1;
		int c = ri() - 1;
		hen[x].push_back({y, c});
		hen[y].push_back({x, c});
	}
	run();
	printf("%" PRId64 "\n", res);
	return 0;
}