結果
問題 |
No.1002 Twotone
|
ユーザー |
![]() |
提出日時 | 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 |
ソースコード
#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; }