cp-library

C++ Library for Competitive Programming

View the Project on GitHub emthrm/cp-library

:heavy_check_mark: グラフ/木/double sweep
(test/graph/tree/double_sweep.test.cpp)

Depends on

Code

/*
 * @title グラフ/木/double sweep
 *
 * verification-helper: PROBLEM https://judge.yosupo.jp/problem/tree_diameter
 */

#include <iostream>
#include <vector>

#include "emthrm/graph/edge.hpp"
#include "emthrm/graph/tree/double_sweep.hpp"

int main() {
  int n;
  std::cin >> n;
  std::vector<std::vector<emthrm::Edge<long long>>> graph(n);
  for (int i = 0; i < n - 1; ++i) {
    int a, b, c;
    std::cin >> a >> b >> c;
    graph[a].emplace_back(a, b, c);
    graph[b].emplace_back(b, a, c);
  }
  const auto [x, u] = emthrm::double_sweep(graph);
  const int y = u.size();
  std::cout << x << ' ' << y << '\n';
  for (int i = 0; i < y; ++i) {
    std::cout << u[i] << " \n"[i + 1 == y];
  }
  return 0;
}
#line 1 "test/graph/tree/double_sweep.test.cpp"
/*
 * @title グラフ/木/double sweep
 *
 * verification-helper: PROBLEM https://judge.yosupo.jp/problem/tree_diameter
 */

#include <iostream>
#include <vector>

#line 1 "include/emthrm/graph/edge.hpp"
/**
 * @title 辺
 */

#ifndef EMTHRM_GRAPH_EDGE_HPP_
#define EMTHRM_GRAPH_EDGE_HPP_

#include <compare>

namespace emthrm {

template <typename CostType>
struct Edge {
  CostType cost;
  int src, dst;

  explicit Edge(const int src, const int dst, const CostType cost = 0)
      : cost(cost), src(src), dst(dst) {}

  auto operator<=>(const Edge& x) const = default;
};

}  // namespace emthrm

#endif  // EMTHRM_GRAPH_EDGE_HPP_
#line 1 "include/emthrm/graph/tree/double_sweep.hpp"



#include <cassert>
#include <ranges>
#include <tuple>
#include <utility>
#line 9 "include/emthrm/graph/tree/double_sweep.hpp"

#line 1 "include/emthrm/graph/edge.hpp"
/**
 * @title 辺
 */

#ifndef EMTHRM_GRAPH_EDGE_HPP_
#define EMTHRM_GRAPH_EDGE_HPP_

#include <compare>

namespace emthrm {

template <typename CostType>
struct Edge {
  CostType cost;
  int src, dst;

  explicit Edge(const int src, const int dst, const CostType cost = 0)
      : cost(cost), src(src), dst(dst) {}

  auto operator<=>(const Edge& x) const = default;
};

}  // namespace emthrm

#endif  // EMTHRM_GRAPH_EDGE_HPP_
#line 11 "include/emthrm/graph/tree/double_sweep.hpp"

namespace emthrm {

template <typename CostType>
std::pair<CostType, std::vector<int>> double_sweep(
    const std::vector<std::vector<Edge<CostType>>>& graph) {
  const auto dfs1 = [&graph](auto dfs1, const int par, const int ver)
      -> std::pair<CostType, int> {
    std::pair<CostType, int> res{0, ver};
    for (const Edge<CostType>& e : graph[ver]) {
      if (e.dst != par) {
        std::pair<CostType, int> child = dfs1(dfs1, ver, e.dst);
        child.first += e.cost;
        if (child.first > res.first) res = child;
      }
    }
    return res;
  };
  const int s = dfs1(dfs1, -1, 0).second;
  const auto [diameter, t] = dfs1(dfs1, -1, s);
  std::vector<int> path{s};
  const auto dfs2 = [&graph, t, &path](auto dfs2, const int par, const int ver)
      -> bool {
    if (ver == t) return true;
    for (const int e : graph[ver]
                     | std::views::transform(&Edge<CostType>::dst)) {
      if (e != par) {
        path.emplace_back(e);
        if (dfs2(dfs2, ver, e)) return true;
        path.pop_back();
      }
    }
    return false;
  };
  assert(dfs2(dfs2, -1, s));
  return {diameter, path};
}

}  // namespace emthrm


#line 12 "test/graph/tree/double_sweep.test.cpp"

int main() {
  int n;
  std::cin >> n;
  std::vector<std::vector<emthrm::Edge<long long>>> graph(n);
  for (int i = 0; i < n - 1; ++i) {
    int a, b, c;
    std::cin >> a >> b >> c;
    graph[a].emplace_back(a, b, c);
    graph[b].emplace_back(b, a, c);
  }
  const auto [x, u] = emthrm::double_sweep(graph);
  const int y = u.size();
  std::cout << x << ' ' << y << '\n';
  for (int i = 0; i < y; ++i) {
    std::cout << u[i] << " \n"[i + 1 == y];
  }
  return 0;
}
Back to top page