Tree Diameter(木の直径)
(graph/diameter.hpp)

• View this file on GitHub
• Last update: 2024-09-07 20:55:28+09:00
• Include: #include "graph/diameter.hpp"

木の直径を求める。

Depends on

• graph/graph_template.hpp

Verified with

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

Code

#ifndef HARUILIB_GRAPH_DIAMETER_HPP
#define HARUILIB_GRAPH_DIAMETER_HPP


#include <vector>
#include <queue>
#include "graph/graph_template.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

#endif // HARUILIB_GRAPH_DIAMETER_HPP

#line 1 "graph/diameter.hpp"




#include <vector>
#include <queue>
#line 1 "graph/graph_template.hpp"



#line 5 "graph/graph_template.hpp"

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 8 "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

Back to top page