結果
問題 | No.1002 Twotone |
ユーザー |
![]() |
提出日時 | 2020-02-28 22:55:45 |
言語 | C++14 (gcc 13.3.0 + boost 1.87.0) |
結果 |
WA
|
実行時間 | - |
コード長 | 2,906 bytes |
コンパイル時間 | 2,555 ms |
コンパイル使用メモリ | 199,252 KB |
実行使用メモリ | 81,496 KB |
最終ジャッジ日時 | 2024-10-13 18:47:54 |
合計ジャッジ時間 | 40,951 ms |
ジャッジサーバーID (参考情報) |
judge4 / judge3 |
(要ログイン)
ファイルパターン | 結果 |
---|---|
sample | AC * 2 |
other | AC * 12 WA * 20 TLE * 1 |
ソースコード
#include <bits/stdc++.h> using namespace std; int main() { cin.tie(nullptr); ios::sync_with_stdio(false); int n, k; cin >> n >> k; vector<vector<int>> g(n); map<pair<int, int>, int> col; for (int _ = n - 1; _--; ) { int u, v, c; cin >> u >> v >> c; --u, --v, --c; g[u].push_back(v); g[v].push_back(u); col[{u, v}] = col[{v, u}] = c; } vector<bool> be(n, true); vector<int> sz(n), par(n); auto centroid = [&](int r) { vector<int> vs; auto dfs = [&](auto&& self, int v, int p) -> void { vs.push_back(v); par[v] = p; sz[v] = 1; for (int u : g[v]) { if (be[u] and u != p) { self(self, u, v); sz[v] += sz[u]; } } }; dfs(dfs, r, -1); for (int v : vs) { bool ok = vs.size() - sz[v] <= vs.size() / 2; for (int u : g[v]) { if (be[u] and u != par[v]) { if (sz[u] > (int)vs.size() / 2) { ok = false; break; } } } if (ok) { return v; } } return -1; }; long long res = 0; map<int, int> mp1; map<pair<int, int>, int> mp2; auto dfs = [&](auto&& self, int v, int p, set<int> se) -> void { if (se.size() == 1) { ++mp1[*begin(se)]; } else if (se.size() == 2) { ++mp2[{*begin(se), *next(begin(se))}]; } else { assert(false); } for (int u : g[v]) { if (be[u] and u != p) { auto nse = se; nse.insert(col[{v, u}]); if (nse.size() >= 3) { continue; } self(self, u, v, nse); } } }; auto rec = [&](auto&& self, int c) -> void { c = centroid(c); map<int, int> smp1; map<pair<int, int>, int> smp2; for (int v : g[c]) { if (be[v]) { mp1.clear(), mp2.clear(); dfs(dfs, v, c, {col[{c, v}]}); for (auto e : mp1) { for (auto f : mp1) { if (e.first != f.first) { res -= (long long)e.second * f.second; } } } for (auto e : mp2) { res += 2 * e.second; res -= (long long)(mp1[e.first.first] + mp1[e.first.second]) * e.second; res -= (long long)e.second * e.second; } for (auto e : mp1) { smp1[e.first] += e.second; } for (auto e : mp2) { smp2[e.first] += e.second; } } } for (auto e : smp1) { for (auto f : smp1) { if (e.first != f.first) { res += (long long)e.second * f.second; } } } for (auto e : smp2) { res += (long long)(smp1[e.first.first] + smp1[e.first.second]) * e.second; res += (long long)e.second * e.second; } be[c] = false; for (int u : g[c]) { if (be[u]) { self(self, u); } } be[c] = true; }; rec(rec, 0); res /= 2; cout << res << '\n'; }