sotanishy's code snippets for competitive programming
View the Project on GitHub sotanishy/cp-library-cpp
#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter" #include "../../tree/tree_diameter.hpp" #include <bits/stdc++.h> using namespace std; using ll = long long; template <typename T> ostream& operator<<(ostream& os, const vector<T>& v) { for (int i = 0; i < (int)v.size(); ++i) cout << v[i] << " "; return os; } int main() { ios_base::sync_with_stdio(false); cin.tie(nullptr); int N; cin >> N; vector<vector<pair<int, ll>>> G(N); for (int i = 0; i < N - 1; ++i) { int a, b, c; cin >> a >> b >> c; G[a].push_back({b, c}); G[b].push_back({a, c}); } auto ans = tree_diameter(G); cout << ans.first << " " << ans.second.size() << endl; cout << ans.second << endl; }
#line 1 "test/yosupo/tree_diameter.test.cpp" #define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter" #line 2 "tree/tree_diameter.hpp" #include <algorithm> #include <utility> #include <vector> std::pair<int, std::vector<int>> tree_diameter( const std::vector<std::vector<int>>& G) { std::vector<int> to(G.size()); auto dfs = [&](const auto& dfs, int v, int p) -> std::pair<int, int> { std::pair<int, int> ret(0, v); for (int c : G[v]) { if (c == p) continue; auto weight = dfs(dfs, c, v); ++weight.first; if (ret < weight) { ret = weight; to[v] = c; } } return ret; }; auto p = dfs(dfs, 0, -1); auto q = dfs(dfs, p.second, -1); std::vector<int> path; int v = p.second; while (v != q.second) { path.push_back(v); v = to[v]; } path.push_back(v); return {q.first, path}; } template <typename T> std::pair<T, std::vector<int>> tree_diameter( const std::vector<std::vector<std::pair<int, T>>>& G) { std::vector<int> to(G.size()); auto dfs = [&](const auto& dfs, int v, int p) -> std::pair<T, int> { std::pair<T, int> ret(0, v); for (auto& [u, w] : G[v]) { if (u == p) continue; auto weight = dfs(dfs, u, v); weight.first += w; if (ret < weight) { ret = weight; to[v] = u; } } return ret; }; auto p = dfs(dfs, 0, -1); auto q = dfs(dfs, p.second, -1); std::vector<int> path; int v = p.second; while (v != q.second) { path.push_back(v); v = to[v]; } path.push_back(v); return {q.first, path}; } #line 4 "test/yosupo/tree_diameter.test.cpp" #include <bits/stdc++.h> using namespace std; using ll = long long; template <typename T> ostream& operator<<(ostream& os, const vector<T>& v) { for (int i = 0; i < (int)v.size(); ++i) cout << v[i] << " "; return os; } int main() { ios_base::sync_with_stdio(false); cin.tie(nullptr); int N; cin >> N; vector<vector<pair<int, ll>>> G(N); for (int i = 0; i < N - 1; ++i) { int a, b, c; cin >> a >> b >> c; G[a].push_back({b, c}); G[b].push_back({a, c}); } auto ans = tree_diameter(G); cout << ans.first << " " << ans.second.size() << endl; cout << ans.second << endl; }