test/tree_diameter.test.cpp

• View this file on GitHub
• Last update: 2024-02-15 14:50:41+09:00
• Problem: https://judge.yosupo.jp/problem/tree_diameter

Depends on

• tree/tree_diameter.hpp

Code

#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;



int main() {
  ll N;cin>>N;
  vector<vector<pair<int,ll>>> vec(N,vector<pair<int,ll>> (0));
  for(ll i=0;i<N-1;i++) {
    ll a,b,c;cin>>a>>b>>c;
    vec[a].push_back(make_pair(b,c));
    vec[b].push_back(make_pair(a,c));
  }


  auto ans=tree_diameter<ll>(vec);

  cout<<ans.first<<" "<<ans.second.size()<<endl;
  for(auto x:ans.second) cout<<x<<" ";
  cout<<endl;
}

#line 1 "test/tree_diameter.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"

#line 2 "tree/tree_diameter.hpp"
#include<bits/stdc++.h>
using namespace std;

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/tree_diameter.test.cpp"

#line 6 "test/tree_diameter.test.cpp"
using namespace std;
using ll =long long;



int main() {
  ll N;cin>>N;
  vector<vector<pair<int,ll>>> vec(N,vector<pair<int,ll>> (0));
  for(ll i=0;i<N-1;i++) {
    ll a,b,c;cin>>a>>b>>c;
    vec[a].push_back(make_pair(b,c));
    vec[b].push_back(make_pair(a,c));
  }


  auto ans=tree_diameter<ll>(vec);

  cout<<ans.first<<" "<<ans.second.size()<<endl;
  for(auto x:ans.second) cout<<x<<" ";
  cout<<endl;
}

Back to top page