結果

問題 No.1002 Twotone
ユーザー chocorusk
提出日時 2020-02-28 22:49:11
言語 C++14
(gcc 13.3.0 + boost 1.87.0)
結果
AC  
実行時間 1,435 ms / 5,000 ms
コード長 4,023 bytes
コンパイル時間 1,621 ms
コンパイル使用メモリ 136,436 KB
実行使用メモリ 39,680 KB
最終ジャッジ日時 2024-10-13 18:41:39
合計ジャッジ時間 19,053 ms
ジャッジサーバーID
（参考情報） 		judge3 / judge4
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 2
other AC * 33
権限があれば一括ダウンロードができます

ソースコード

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