結果

問題 No.1002 Twotone
ユーザー chocoruskchocorusk
提出日時 2020-02-28 22:49:11
言語 C++14
(gcc 12.3.0 + boost 1.83.0)
結果
AC  
実行時間 1,898 ms / 5,000 ms
コード長 4,023 bytes
コンパイル時間 1,818 ms
コンパイル使用メモリ 135,144 KB
実行使用メモリ 39,808 KB
最終ジャッジ日時 2024-04-21 20:00:03
合計ジャッジ時間 24,027 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge2
このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示
入力 結果 実行時間
実行使用メモリ
testcase_00 AC 5 ms
8,320 KB
testcase_01 AC 6 ms
8,320 KB
testcase_02 AC 6 ms
8,320 KB
testcase_03 AC 342 ms
15,360 KB
testcase_04 AC 471 ms
17,664 KB
testcase_05 AC 444 ms
17,792 KB
testcase_06 AC 6 ms
8,448 KB
testcase_07 AC 435 ms
14,080 KB
testcase_08 AC 794 ms
17,536 KB
testcase_09 AC 799 ms
17,664 KB
testcase_10 AC 6 ms
8,320 KB
testcase_11 AC 1,043 ms
16,128 KB
testcase_12 AC 1,358 ms
18,432 KB
testcase_13 AC 1,350 ms
18,432 KB
testcase_14 AC 5 ms
8,320 KB
testcase_15 AC 656 ms
14,080 KB
testcase_16 AC 1,357 ms
18,176 KB
testcase_17 AC 1,365 ms
18,304 KB
testcase_18 AC 6 ms
8,448 KB
testcase_19 AC 759 ms
24,576 KB
testcase_20 AC 894 ms
35,072 KB
testcase_21 AC 905 ms
34,176 KB
testcase_22 AC 6 ms
8,448 KB
testcase_23 AC 1,008 ms
28,032 KB
testcase_24 AC 1,863 ms
39,808 KB
testcase_25 AC 1,898 ms
39,808 KB
testcase_26 AC 7 ms
8,448 KB
testcase_27 AC 105 ms
17,312 KB
testcase_28 AC 200 ms
20,676 KB
testcase_29 AC 165 ms
20,672 KB
testcase_30 AC 7 ms
8,448 KB
testcase_31 AC 146 ms
20,628 KB
testcase_32 AC 185 ms
20,672 KB
testcase_33 AC 154 ms
20,676 KB
testcase_34 AC 1,170 ms
33,664 KB
権限があれば一括ダウンロードができます

ソースコード

diff # 

#include <cstdio>
#include <cstring>
#include <iostream>
#include <string>
#include <cmath>
#include <bitset>
#include <vector>
#include <map>
#include <set>
#include <queue>
#include <deque>
#include <algorithm>
#include <complex>
#include <unordered_map>
#include <unordered_set>
#include <random>
#include <cassert>
#include <fstream>
#include <utility>
#include <functional>
#include <time.h>
#include <stack>
#include <array>
#define popcount __builtin_popcount
using namespace std;
typedef long long int ll;
typedef pair<int, int> P;
const int MAX_V = 200020;  // ツリーのサイズのありうる最大値

int N;  // ツリーのサイズ
vector<P> tree[MAX_V];  // ツリーを隣接リスト形式のグラフ構造で表したもの

int sizeSubtree[MAX_V];  // sizeSubtree[v] := v を根とする部分ツリーのサイズ (分割統治の毎ステップごとに再利用)
bool isRemoved[MAX_V];  // isRemoved[v] := v が既に取り除かれたかどうか
int whoIsParent[MAX_V];  // whoIsParent[v] := ツリーDP時に v の親が誰だったか

// メイン処理
vector<int> centroids;
void FindCentroidRecursive(int v, int size, int p = -1) {
    sizeSubtree[v] = 1;
    whoIsParent[v] = p;
    bool isCentroid = true;
    for (auto pr : tree[v]) {
		int ch=pr.first;
        if (ch == p) continue;
        if (isRemoved[ch]) continue;
        FindCentroidRecursive(ch, size, v);
        if (sizeSubtree[ch] > size / 2) isCentroid = false;
        sizeSubtree[v] += sizeSubtree[ch];
    }
    if (size - sizeSubtree[v] > size / 2) isCentroid = false;
    if (isCentroid) centroids.push_back(v);
}

// 初期化
void Init() {
    for (int i = 0; i < MAX_V; ++i) {
        isRemoved[i] = false;
    }
}

// first: 重心, second: (重心の子ノード, 子部分木のサイズ) からなるベクトル
pair<int, vector<pair<int,int> > > FindCentroid(int root, int size) {
    vector<pair<int, int> > subtrees;
    centroids.clear();
    FindCentroidRecursive(root, size);
    int center = centroids[0];
    for (auto pr : tree[center]) {
		int ch=pr.first;
        if (isRemoved[ch]) continue;
        if (ch == whoIsParent[center]) {
            subtrees.push_back(make_pair(ch, size - sizeSubtree[center]));
        }
        else {
            subtrees.push_back(make_pair(ch, sizeSubtree[ch]));
        }
    }
    return make_pair(center, subtrees);
}
map<P, ll> mp1;
map<int, ll> mp2;
void dfs(int v, int p, vector<int> vc){
	//whoIsParent[v] = p;
	if(vc.size()==1) mp2[vc[0]]++;
	else if(vc.size()==2) mp1[{vc[0], vc[1]}]++;
	for(auto pr:tree[v]){
        int ch=pr.first;
        if (ch == p) continue;
        if (isRemoved[ch]) continue;
		vector<int> vc2=vc;
		vc2.push_back(pr.second);
		sort(vc2.begin(), vc2.end());
		vc2.erase(unique(vc2.begin(), vc2.end()), vc2.end());
		if(vc2.size()<=2) dfs(ch, v, vc2);
    }
}
ll ans;
void solve(int root, int size){
	if(size<=1) return;
	pair<int, vector<pair<int,int> > > pr=FindCentroid(root, size);
    int cent=pr.first;
	mp1.clear(); mp2.clear();
	vector<int> vc0;
	dfs(cent, -1, vc0);
	for(auto p:mp1) ans+=p.second;
	auto calc=[&](){
		ll ans1=0;
		for(auto p:mp1){
			ans1+=p.second*p.second;
			int x=p.first.first, y=p.first.second;
			if(mp2.find(x)!=mp2.end()){
				ans1+=p.second*mp2[x]*2;
			}
			if(mp2.find(y)!=mp2.end()){
				ans1+=p.second*mp2[y]*2;
			}
		}
		ll sum=0;
		for(auto p:mp2) sum+=p.second;
		for(auto p:mp2) ans1+=p.second*(sum-p.second);
		return ans1;
	};
	ll ans1=calc();
	isRemoved[cent]=1;
	for(auto pr:tree[cent]){
		int y=pr.first;
		if(!isRemoved[y]){
			mp1.clear(); mp2.clear();
			vector<int> vc0(1, pr.second);
			dfs(y, -1, vc0);
			ans1-=calc();
		}
	}
	ans+=ans1/2;
	for(auto prr:pr.second){
        if(!isRemoved[prr.first]) solve(prr.first, prr.second);
    }
}
int main()
{
    int n, k;
	cin>>n>>k;
	for(int i=0; i<n-1; i++){
		int u, v, c;
		scanf("%d %d %d", &u, &v, &c);
		u--; v--;
		tree[u].push_back({v, c});
		tree[v].push_back({u, c});
	}
	solve(0, n);
	cout<<ans<<endl;
	return 0;
}
0