test/verify/yosupo-tree-diameter.test.cpp

View this file on GitHub
Last update: 2025-03-23 00:55:59+09:00
Problem: https://judge.yosupo.jp/problem/tree_diameter

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/tree_diameter"

#include "template/template.hpp"
#include "graph/graph_template.hpp"
#include "graph/diameter.hpp"

int main() {
  ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);

  int N; cin >> N;
  Graph<ll>graph(N);
  for(int i=0; i<N-1; i++) {
    ll a, b, c; cin >> a >> b >> c;
    graph.add_edge(a,b,c);
    graph.add_edge(b,a,c);
  }

  pair<vector<int>, ll> pr = mylib::diameter_path<ll>(graph);

  cout << pr.second << " " << pr.first.size() << "\n";
  for(int i=0; i<(int)pr.first.size(); i++) {
    cout << pr.first[i];
    if (i < N-1) cout << " ";
    else cout << "\n";
  }

  return 0;
}

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

#line 1 "template/template.hpp"
#include <iostream>
#include <cassert>
using namespace std;
using ll = long long;
template<class T> inline bool chmax(T& a, const T& b) {if (a<b) {a=b; return true;} return false;}
template<class T> inline bool chmin(T& a, const T& b) {if (b<a) {a=b; return true;} return false;}
const int INTINF = 1000001000;
const int INTMAX = 2147483647;
const ll LLMAX = 9223372036854775807;
const ll LLINF = 1000000000000000000;
#line 1 "graph/graph_template.hpp"



#include <vector>

template <typename T>
struct Edge {
  int from; int to;
  T cost;

  // default constructor
  Edge () : from(-1), to(-1), cost(T(0)) {}

  Edge(int _from, int _to, T _cost) : from(_from), to(_to), cost(_cost) {}

  // unweighted
  Edge(int _from, int _to) : from(_from), to(_to), cost(T(1)) {}

  bool operator==(const Edge& rhs) const {
    return from == rhs.from && to == rhs.to && cost == rhs.cost;
  }

  bool operator<(const Edge& rhs) const {
    return cost < rhs.cost;
  }
  
  bool operator>(const Edge& rhs) const {
    return cost > rhs.cost;
  }

};


template <typename T>
struct Graph : std::vector<std::vector<Edge<T>>> {

  using std::vector<std::vector<Edge<T>>>::vector; // inherit constructors

  void add_edge(int i, Edge<T> e) {
    (*this)[i].push_back(e);
  }

  void add_edge(Edge<T> e) {
    (*this)[e.from].push_back(e);
  }

  // weighted
  void add_edge(int _from, int _to, T _cost) {
    (*this)[_from].push_back(Edge(_from, _to, _cost));
  }

  // unweighted
  void add_edge(int _from, int _to) {
    (*this)[_from].push_back(Edge(_from, _to, T(1)));
  }

};


#line 1 "graph/diameter.hpp"




#include <algorithm>
#line 7 "graph/diameter.hpp"
#include <queue>
#line 9 "graph/diameter.hpp"

namespace mylib {

  using std::vector;
  using std::pair;
  using std::queue;

  template <typename T>
  pair<vector<int>, T> diameter_path(const Graph<T>& graph) {
    int N = (int)graph.size();

    // first sweep
    pair<T, int> farest1 = make_pair(T(0), 0);
    {
      vector<T> dist(N, T(-1));
      dist[0] = T(0);

      queue<int> que;
      que.push(0);
      while (!que.empty()) {
        int front = que.front(); que.pop();

        for (Edge e: graph[front]) {
          if (dist[e.to] == T(-1)) {
            dist[e.to] = dist[front] + e.cost;
            que.push(e.to);
            if (farest1.first == -1 || farest1.first < dist[e.to]) {
              farest1 = make_pair(dist[e.to], e.to);
            }
          }
        }
      }
    }


    pair<T, int> farest2 = make_pair(T(0), farest1.second);
    // second sweep
    vector<int> prev(N, -1);
    {
      int start = farest1.second;
      vector<T> dist(N, T(-1));
      dist[start] = T(0);

      queue<int> que;
      que.push(start);

      while (!que.empty()) {
        int front = que.front(); que.pop();

        for (Edge e: graph[front]) {
          if (dist[e.to] == T(-1)) {
            dist[e.to] = dist[front] + e.cost;
            prev[e.to] = front;
            que.push(e.to);
          }

          if (farest2.second == -1 || farest2.first < dist[e.to]) {
            farest2 = make_pair(dist[e.to], e.to);
          }
        }
      }
    }

    vector<int> diameter_path;
    {
      int now = farest2.second;
      while (now != -1) {
        diameter_path.push_back(now);
        now = prev[now];
      }

      reverse(diameter_path.begin(), diameter_path.end());
    }

    return make_pair(diameter_path, farest2.first);
  }

} // namespace mylib


#line 6 "test/verify/yosupo-tree-diameter.test.cpp"

int main() {
  ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);

  int N; cin >> N;
  Graph<ll>graph(N);
  for(int i=0; i<N-1; i++) {
    ll a, b, c; cin >> a >> b >> c;
    graph.add_edge(a,b,c);
    graph.add_edge(b,a,c);
  }

  pair<vector<int>, ll> pr = mylib::diameter_path<ll>(graph);

  cout << pr.second << " " << pr.first.size() << "\n";
  for(int i=0; i<(int)pr.first.size(); i++) {
    cout << pr.first[i];
    if (i < N-1) cout << " ";
    else cout << "\n";
  }

  return 0;
}