tests/prim/basic.test.cpp

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

Depends on

• std::greater by key (tools/greater_by.hpp)
• Join elements with delimiter (tools/join.hpp)
• Prim's algorithm (tools/prim.hpp)

Code

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

#include <iostream>
#include "tools/prim.hpp"
#include "tools/join.hpp"

using ll = long long;

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

  int N, M;
  std::cin >> N >> M;
  tools::prim<ll> graph(N);
  for (int i = 0; i < M; ++i) {
    int a, b, c;
    std::cin >> a >> b >> c;
    graph.add_edge(a, b, c);
  }

  const auto [X, e] = graph.query().first[0];
  std::cout << X << '\n';
  std::cout << tools::join(e, " ") << '\n';

  return 0;
}

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

#include <iostream>
#line 1 "tools/prim.hpp"



#include <algorithm>
#include <cassert>
#include <cstddef>
#include <limits>
#include <queue>
#include <tuple>
#include <utility>
#include <vector>
#line 1 "tools/greater_by.hpp"



namespace tools {

  template <class F>
  class greater_by {
  private:
    F selector;

  public:
    greater_by(const F& selector) : selector(selector) {
    }

    template <class T>
    bool operator()(const T& x, const T& y) const {
      return selector(x) > selector(y);
    }
  };
}


#line 13 "tools/prim.hpp"

namespace tools {

  template <typename T>
  class prim {
  public:
    struct edge {
      ::std::size_t id;
      ::std::size_t from;
      ::std::size_t to;
      T cost;
    };

  private:
    ::std::vector<edge> m_edges;
    ::std::vector<::std::vector<::std::size_t>> m_graph;

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

    explicit prim(const ::std::size_t n) : m_graph(n) {
    }

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

    ::std::size_t add_edge(::std::size_t u, ::std::size_t v, const T w) {
      assert(u < this->size());
      assert(v < this->size());
      ::std::tie(u, v) = ::std::minmax({u, v});
      this->m_edges.push_back(edge({this->m_edges.size(), u, v, w}));
      this->m_graph[u].push_back(this->m_edges.size() - 1);
      this->m_graph[v].push_back(this->m_edges.size() - 1);
      return this->m_edges.size() - 1;
    }

    const edge& get_edge(const ::std::size_t k) const {
      assert(k < this->m_edges.size());
      return this->m_edges[k];
    }

    const ::std::vector<edge>& edges() const {
      return this->m_edges;
    }

    ::std::pair<::std::vector<::std::pair<T, ::std::vector<::std::size_t>>>, ::std::vector<::std::size_t>> query() const {
      ::std::pair<::std::vector<::std::pair<T, ::std::vector<::std::size_t>>>, ::std::vector<::std::size_t>> res;
      auto& [groups, belongs_to] = res;
      belongs_to.resize(this->size());
      ::std::fill(belongs_to.begin(), belongs_to.end(), ::std::numeric_limits<::std::size_t>::max());

      for (::std::size_t root = 0; root < this->size(); ++root) {
        if (belongs_to[root] < ::std::numeric_limits<::std::size_t>::max()) continue;

        const auto greater_by_cost = ::tools::greater_by([&](const auto& pair) { return this->m_edges[pair.first].cost; });
        ::std::priority_queue<::std::pair<::std::size_t, ::std::size_t>, ::std::vector<::std::pair<::std::size_t, ::std::size_t>>, decltype(greater_by_cost)> pq(greater_by_cost);
        groups.emplace_back(0, ::std::vector<::std::size_t>());
        auto& [total_cost, edge_ids] = groups.back();

        belongs_to[root] = groups.size() - 1;
        for (const auto e : this->m_graph[root]) {
          const auto next = this->m_edges[e].from ^ this->m_edges[e].to ^ root;
          if (belongs_to[next] < ::std::numeric_limits<::std::size_t>::max()) continue;

          pq.emplace(e, next);
        }

        while (!pq.empty()) {
          const auto [from_edge_id, here] = pq.top();
          pq.pop();

          if (belongs_to[here] < ::std::numeric_limits<::std::size_t>::max()) continue;

          belongs_to[here] = belongs_to[root];
          total_cost += this->m_edges[from_edge_id].cost;
          edge_ids.push_back(from_edge_id);
          for (const auto e : this->m_graph[here]) {
            const auto next = this->m_edges[e].from ^ this->m_edges[e].to ^ here;
            if (belongs_to[next] < ::std::numeric_limits<::std::size_t>::max()) continue;

            pq.emplace(e, next);
          }
        }
      }

      return res;
    }
  };
}


#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 6 "tests/prim/basic.test.cpp"

using ll = long long;

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

  int N, M;
  std::cin >> N >> M;
  tools::prim<ll> graph(N);
  for (int i = 0; i < M; ++i) {
    int a, b, c;
    std::cin >> a >> b >> c;
    graph.add_edge(a, b, c);
  }

  const auto [X, e] = graph.query().first[0];
  std::cout << X << '\n';
  std::cout << tools::join(e, " ") << '\n';

  return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	4 ms	4 MB
g++	example_01	AC	4 ms	4 MB
g++	example_02	AC	4 ms	4 MB
g++	max_random_00	AC	450 ms	63 MB
g++	max_random_01	AC	418 ms	63 MB
g++	max_random_02	AC	399 ms	63 MB
g++	max_random_03	AC	420 ms	63 MB
g++	max_random_04	AC	435 ms	63 MB
g++	random_00	AC	358 ms	53 MB
g++	random_01	AC	416 ms	63 MB
g++	random_02	AC	325 ms	47 MB
g++	random_03	AC	391 ms	59 MB
g++	random_04	AC	329 ms	48 MB
g++	small_00	AC	263 ms	43 MB
g++	small_01	AC	324 ms	43 MB
g++	small_02	AC	218 ms	42 MB
g++	small_03	AC	16 ms	6 MB
g++	small_04	AC	51 ms	13 MB
g++	small_c01_00	AC	5 ms	4 MB
g++	small_c01_01	AC	4 ms	4 MB
g++	small_c01_02	AC	4 ms	4 MB
g++	small_c01_03	AC	4 ms	4 MB
g++	small_c01_04	AC	4 ms	4 MB
g++	star_00	AC	470 ms	70 MB

Back to top page