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

• 問題: No.1002 Twotone
• ユーザー: QCFium
• 提出日時: 2020-03-02 23:11:41
• 言語: C++14 (gcc 12.3.0 + boost 1.83.0)
• 結果: AC
• 実行時間: 1,208 ms / 5,000 ms
• コード長: 3,590 bytes
• コンパイル時間: 2,735 ms
• コンパイル使用メモリ: 197,856 KB
• 実行使用メモリ: 42,856 KB
• 最終ジャッジ日時: 2024-10-13 21:22:00
• 合計ジャッジ時間: 17,066 ms
• ジャッジサーバーID （参考情報）: judge2 / judge3

入力	結果	実行時間; 実行使用メモリ
testcase_00	AC	1 ms; 5,248 KB
testcase_01	AC	2 ms; 5,248 KB
testcase_02	AC	2 ms; 5,248 KB
testcase_03	AC	271 ms; 14,592 KB
testcase_04	AC	356 ms; 17,920 KB
testcase_05	AC	357 ms; 17,920 KB
testcase_06	AC	2 ms; 5,248 KB
testcase_07	AC	256 ms; 16,728 KB
testcase_08	AC	455 ms; 24,952 KB
testcase_09	AC	459 ms; 25,092 KB
testcase_10	AC	2 ms; 5,248 KB
testcase_11	AC	494 ms; 18,312 KB
testcase_12	AC	681 ms; 22,784 KB
testcase_13	AC	703 ms; 22,920 KB
testcase_14	AC	3 ms; 5,248 KB
testcase_15	AC	450 ms; 18,324 KB
testcase_16	AC	912 ms; 29,260 KB
testcase_17	AC	877 ms; 27,740 KB
testcase_18	AC	2 ms; 5,248 KB
testcase_19	AC	532 ms; 22,912 KB
testcase_20	AC	660 ms; 32,256 KB
testcase_21	AC	638 ms; 31,616 KB
testcase_22	AC	2 ms; 5,248 KB
testcase_23	AC	662 ms; 26,080 KB
testcase_24	AC	1,191 ms; 42,324 KB
testcase_25	AC	1,208 ms; 40,252 KB
testcase_26	AC	2 ms; 5,248 KB
testcase_27	AC	122 ms; 37,696 KB
testcase_28	AC	199 ms; 42,732 KB
testcase_29	AC	186 ms; 42,852 KB
testcase_30	AC	2 ms; 5,248 KB
testcase_31	AC	177 ms; 42,664 KB
testcase_32	AC	202 ms; 42,724 KB
testcase_33	AC	188 ms; 42,856 KB
testcase_34	AC	942 ms; 34,912 KB

ソースコード

#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);
	int init_size = size[i];
	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 > init_size / 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::vector<std::pair<int, int> > two_match;
		std::vector<int> two_side;
		std::vector<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.push_back(cur_color);
				info.two_side.push_back(cur_color.first);
				info.two_side.push_back(cur_color.second);
			} else {
				info.one_num++;
				info.one.push_back(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});
			}
		}
		std::sort(info.two_match.begin(), info.two_match.end());
		std::sort(info.two_side.begin(), info.two_side.end());
		std::sort(info.one.begin(), info.one.end());
	}
	Info all;
	for (auto &j : infos) {
		for (auto &k : j.two_match) all.two_match.push_back(k);
		for (auto &k : j.two_side) all.two_side.push_back(k);
		for (auto &k : j.one) all.one.push_back(k);
		all.one_num += j.one_num;
	}
	std::sort(all.two_match.begin(), all.two_match.end());
	std::sort(all.two_side.begin(), all.two_side.end());
	std::sort(all.one.begin(), all.one.end());
	auto count_in = [&] (auto &array, auto val) {
		return std::upper_bound(array.begin(), array.end(), val) - std::lower_bound(array.begin(), array.end(), val);
	};
	auto calc = [&] (Info &target, Info &main) {
		int64_t res = 0;
		for (auto &j : main.two_match) {
			res += count_in(target.two_match, j);
			res += count_in(target.one, j.first) + count_in(target.one, j.second);
		}
		for (auto &j : main.one) {
			res += count_in(target.two_side, j);
			res += target.one_num - count_in(target.one, j);
		}
		return res;
	};
	for (auto &j : infos) res += calc(all, j) - calc(j, j);
	res += all.two_match.size() * 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;
}