Prim's algorithm (tools/prim.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2025-02-01 15:01:54+09:00
• Include: #include "tools/prim.hpp"

It constructs minimum spanning trees for each connected components of a given undirected graph which is not necessarily simple based on Prim’s algorithm.

License

• CC0

Author

• anqooqie

Constructor

prim<T> graph(std::size_t n);

It creates an undirected graph with $n$ vertices and $0$ edges. The type parameter <T> represents the type of the cost.

Constraints

• None

Time Complexity

• $O(n)$

size

std::size_t graph.size();

It returns $n$.

Constraints

• None

Time Complexity

• $O(1)$

add_edge

std::size_t graph.add_edge(std::size_t u, std::size_t v, T w);

It adds an undirected edge between $u$ and $v$ with cost w. It returns an integer $k$ such that this is the $k$-th ($0$ indexed) edge that is added.

Constraints

• $0 \leq u < n$
• $0 \leq v < n$

Time Complexity

• $O(1)$ amortized

get_edge

struct edge {
  std::size_t id;
  std::size_t from;
  std::size_t to;
  T cost;
};
const edge& graph.get_edge(std::size_t k);

It returns the $k$-th ($0$ indexed) edge.

Constraints

• $0 \leq k < |E|$ where $|E|$ is the number of edges

Time Complexity

• $O(1)$

edges

const std::vector<edge>& graph.edges();

It returns all the edges in the graph. The edges are ordered in the same order as added by add_edge.

Constraints

• None

Time Complexity

• $O(1)$

query

std::pair<std::vector<std::pair<T, std::vector<std::size_t>>>, std::vector<std::size_t>> graph.query();

It constructs minimum spanning trees for each connected components of the graph, and returns the information about the minimum spanning trees in the following layout. The order of minimum spanning trees is undefined.

(
  [
    (
      total cost of the 0th MST,
      [
        index of an edge in the 0th MST,
        index of another edge in the 0th MST,
        ...
      ]
    ),
    (
      total cost of the 1st MST,
      [
        index of an edge in the 1st MST,
        index of another edge in the 1st MST,
        ...
      ]
    ),
    (
      total cost of the 2nd MST,
      [
        index of an edge in the 2nd MST,
        index of another edge in the 2nd MST,
        ...
      ]
    ),
    ...
  ],
  [
    index of the MST which contains the 0th vertex,
    index of the MST which contains the 1st vertex,
    index of the MST which contains the 2nd vertex,
    ...
  ]
)

Constraints

• None

Time Complexity

• $O((n + |E|) \log n)$ where $|E|$ is the number of edges

Depends on

• std::greater by key (tools/greater_by.hpp)

Verified with

• tests/prim/basic.test.cpp
• tests/prim/unconnected.test.cpp

Code

#ifndef TOOLS_PRIM_HPP
#define TOOLS_PRIM_HPP

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <limits>
#include <queue>
#include <tuple>
#include <utility>
#include <vector>
#include "tools/greater_by.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;
    }
  };
}

#endif

#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;
    }
  };
}

Back to top page