結果

問題 No.1002 Twotone
ユーザー QCFiumQCFium
提出日時 2020-03-02 22:28:08
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
TLE  
実行時間 -
コード長 3,165 bytes
コンパイル時間 2,791 ms
コンパイル使用メモリ 203,752 KB
実行使用メモリ 33,664 KB
最終ジャッジ日時 2024-10-13 21:20:13
合計ジャッジ時間 14,838 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge4
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 2 ms
5,248 KB
testcase_01 AC 2 ms
5,248 KB
testcase_02 AC 1 ms
5,248 KB
testcase_03 AC 348 ms
14,720 KB
testcase_04 AC 489 ms
18,240 KB
testcase_05 AC 483 ms
18,304 KB
testcase_06 AC 2 ms
5,248 KB
testcase_07 AC 188 ms
12,288 KB
testcase_08 AC 339 ms
17,920 KB
testcase_09 AC 343 ms
18,024 KB
testcase_10 AC 2 ms
5,248 KB
testcase_11 AC 303 ms
16,128 KB
testcase_12 AC 468 ms
20,088 KB
testcase_13 AC 488 ms
20,096 KB
testcase_14 AC 2 ms
5,248 KB
testcase_15 AC 222 ms
12,416 KB
testcase_16 AC 446 ms
18,816 KB
testcase_17 AC 480 ms
18,980 KB
testcase_18 AC 2 ms
5,248 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 -- -
権限があれば一括ダウンロードができます

ソースコード

diff # 

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