tests/tree_diameter.test.cpp

• View this file on GitHub
• Last update: 2025-05-10 19:26:48+09:00
• Problem: https://judge.yosupo.jp/problem/tree_diameter

Depends on

• chmin function (tools/chmin.hpp)
• Join elements with delimiter (tools/join.hpp)
• Diameter of a tree (tools/tree_diameter.hpp)

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/tree_diameter

#include <iostream>
#include <vector>
#include <utility>
#include "tools/tree_diameter.hpp"
#include "tools/join.hpp"

using ll = long long;

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  ll N;
  std::cin >> N;
  tools::tree_diameter<ll> tree(N);
  std::vector<std::pair<ll, ll>> edges;
  for (ll i = 0; i < N - 1; ++i) {
    ll a, b, c;
    std::cin >> a >> b >> c;
    tree.add_edge(a, b, c);
    edges.emplace_back(a, b);
  }

  const auto [X, u, unused] = tree.query();
  std::cout << X << ' ' << u.size() << '\n';
  std::cout << ::tools::join(u, " ") << '\n';
  return 0;
}

#line 1 "tests/tree_diameter.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/tree_diameter

#include <iostream>
#include <vector>
#include <utility>
#line 1 "tools/tree_diameter.hpp"



#line 5 "tools/tree_diameter.hpp"
#include <cstddef>
#line 7 "tools/tree_diameter.hpp"
#include <cassert>
#include <tuple>
#include <limits>
#include <queue>
#include <iterator>
#include <algorithm>
#line 1 "tools/chmin.hpp"



#include <type_traits>
#line 6 "tools/chmin.hpp"

namespace tools {

  template <typename M, typename N>
  bool chmin(M& lhs, const N& rhs) {
    bool updated;
    if constexpr (::std::is_integral_v<M> && ::std::is_integral_v<N>) {
      updated = ::std::cmp_less(rhs, lhs);
    } else {
      updated = rhs < lhs;
    }
    if (updated) lhs = rhs;
    return updated;
  }
}


#line 14 "tools/tree_diameter.hpp"

namespace tools {
  template <typename T>
  class tree_diameter {
  private:
    ::std::vector<::std::vector<::std::size_t>> m_graph;
    ::std::vector<::std::pair<::std::size_t, T>> m_edges;

  public:
    tree_diameter() = default;
    tree_diameter(const ::tools::tree_diameter<T>&) = default;
    tree_diameter(::tools::tree_diameter<T>&&) = default;
    ~tree_diameter() = default;
    ::tools::tree_diameter<T>& operator=(const ::tools::tree_diameter<T>&) = default;
    ::tools::tree_diameter<T>& operator=(::tools::tree_diameter<T>&&) = default;

    explicit tree_diameter(const ::std::size_t n) :
      m_graph(n) {
      assert(n >= 1);
    }

    ::std::size_t size() const {
      return this->m_graph.size();
    }

    ::std::size_t add_edge(const ::std::size_t u, const ::std::size_t v, const T& w) {
      assert(u < this->size());
      assert(v < this->size());

      this->m_graph[u].push_back(this->m_edges.size());
      this->m_graph[v].push_back(this->m_edges.size());
      this->m_edges.emplace_back(u ^ v, w);
      return this->m_edges.size() - 1;
    }

    ::std::tuple<T, ::std::vector<::std::size_t>, ::std::vector<::std::size_t>> query() const {
      assert(this->m_edges.size() + 1 == this->size());

      ::std::vector<T> dist(this->size(), ::std::numeric_limits<T>::max());
      dist[0] = 0;
      ::std::queue<::std::size_t> queue({0});
      while (!queue.empty()) {
        const auto here = queue.front();
        queue.pop();
        for (const auto eid : this->m_graph[here]) {
          const auto next = this->m_edges[eid].first ^ here;
          const auto w = this->m_edges[eid].second;
          if (::tools::chmin(dist[next], dist[here] + w)) {
            queue.push(next);
          }
        }
      }

      queue.push(::std::distance(dist.begin(), ::std::max_element(dist.begin(), dist.end())));
      ::std::fill(dist.begin(), dist.end(), ::std::numeric_limits<T>::max());
      dist[queue.front()] = 0;
      ::std::vector<::std::size_t> prev(this->size(), ::std::numeric_limits<::std::size_t>::max());
      while (!queue.empty()) {
        const auto here = queue.front();
        queue.pop();
        for (const auto eid : this->m_graph[here]) {
          const auto next = this->m_edges[eid].first ^ here;
          const auto w = this->m_edges[eid].second;
          if (::tools::chmin(dist[next], dist[here] + w)) {
            prev[next] = eid;
            queue.push(next);
          }
        }
      }

      ::std::tuple<T, ::std::vector<::std::size_t>, ::std::vector<::std::size_t>> result;
      ::std::get<0>(result) = 0;
      ::std::size_t v;
      for (v = ::std::distance(dist.begin(), ::std::max_element(dist.begin(), dist.end())); prev[v] != ::std::numeric_limits<::std::size_t>::max(); v = this->m_edges[prev[v]].first ^ v) {
        ::std::get<0>(result) += this->m_edges[prev[v]].second;
        ::std::get<1>(result).push_back(v);
        ::std::get<2>(result).push_back(prev[v]);
      }
      ::std::get<1>(result).push_back(v);
      return result;
    }
  };
}


#line 1 "tools/join.hpp"



#include <ranges>
#include <sstream>

namespace tools {
  template <::std::ranges::range R, typename T>
  ::std::string join(R&& e, const T& d) {
    ::std::ostringstream ss;
    auto it = ::std::ranges::begin(e);
    const auto end = ::std::ranges::end(e);
    if (it != end) {
      ss << *it;
      for (++it; it != end; ++it) {
        ss << d << *it;
      }
    }
    return ss.str();
  }
}


#line 8 "tests/tree_diameter.test.cpp"

using ll = long long;

int main() {
  std::cin.tie(nullptr);
  std::ios_base::sync_with_stdio(false);

  ll N;
  std::cin >> N;
  tools::tree_diameter<ll> tree(N);
  std::vector<std::pair<ll, ll>> edges;
  for (ll i = 0; i < N - 1; ++i) {
    ll a, b, c;
    std::cin >> a >> b >> c;
    tree.add_edge(a, b, c);
    edges.emplace_back(a, b);
  }

  const auto [X, u, unused] = tree.query();
  std::cout << X << ' ' << u.size() << '\n';
  std::cout << ::tools::join(u, " ") << '\n';
  return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	5 ms	4 MB
g++	hack_00	AC	4 ms	4 MB
g++	line_00	AC	335 ms	68 MB
g++	max_random_00	AC	280 ms	58 MB
g++	max_random_01	AC	273 ms	58 MB
g++	random_00	AC	193 ms	49 MB
g++	random_01	AC	261 ms	54 MB
g++	random_02	AC	25 ms	9 MB
g++	random_03	AC	229 ms	52 MB
g++	random_04	AC	123 ms	35 MB
g++	small_random_00	AC	5 ms	4 MB
g++	small_random_01	AC	4 ms	4 MB
g++	small_random_02	AC	4 ms	4 MB
g++	small_random_03	AC	4 ms	4 MB
g++	small_random_04	AC	4 ms	4 MB
g++	uni_00	AC	178 ms	62 MB

Back to top page