結果

問題 No.1002 Twotone
コンテスト
ユーザー QCFium
提出日時 2020-03-02 22:51:01
言語 C++14
(gcc 13.3.0 + boost 1.87.0)
結果
TLE  
実行時間 -
コード長 3,589 bytes
コンパイル時間 2,357 ms
コンパイル使用メモリ 199,468 KB
実行使用メモリ 21,464 KB
最終ジャッジ日時 2024-10-13 21:21:06
合計ジャッジ時間 10,685 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 5 TLE * 1 -- * 27
権限があれば一括ダウンロードができます

ソースコード

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