Implement undirected graph

Author: constantinpape

CMakeLists.txt

@@ -12,6 +12,7 @@ find_package(nanobind CONFIG REQUIRED)
 nanobind_add_module(_core
     NB_STATIC
     src/bindings/module.cxx
+    src/bindings/graph.cxx
     src/bindings/segmentation.cxx
     src/bindings/utils.cxx
     src/cpp/segmentation/mutex_watershed.cxx

README.md

@@ -41,6 +41,12 @@ segmentation = bic.segmentation.mutex_watershed(
 )
 ```
 
+```python
+graph = bic.graph.UndirectedGraph.from_edges(4, [[0, 1], [1, 2], [2, 3]])
+graph.find_edge(2, 1)
+# 1
+```
+
 ## Scope
 
 The project is not a compatibility layer for `nifty`, `vigra`, or other large libraries. It keeps I/O and heavy dependencies out of the C++ core; callers should use existing Python packages for file formats and pass NumPy arrays into `bioimage-cpp`.

include/bioimage_cpp/graph/undirected_graph.hxx

@@ -0,0 +1,217 @@
+#pragma once
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+namespace bioimage_cpp::graph {
+
+struct Adjacency {
+    std::uint64_t node;
+    std::uint64_t edge;
+};
+
+class UndirectedGraph {
+public:
+    using NodeId = std::uint64_t;
+    using EdgeId = std::uint64_t;
+    using Edge = std::pair<NodeId, NodeId>;
+    using AdjacencyList = std::vector<Adjacency>;
+
+    explicit UndirectedGraph(
+        const NodeId number_of_nodes = 0,
+        const EdgeId reserve_number_of_edges = 0
+    )
+        : number_of_nodes_(number_of_nodes),
+          adjacency_(static_cast<std::size_t>(number_of_nodes)) {
+        edges_.reserve(static_cast<std::size_t>(reserve_number_of_edges));
+    }
+
+    virtual ~UndirectedGraph() = default;
+
+    void assign(
+        const NodeId number_of_nodes = 0,
+        const EdgeId reserve_number_of_edges = 0
+    ) {
+        number_of_nodes_ = number_of_nodes;
+        edges_.clear();
+        edge_lookup_.clear();
+        adjacency_.clear();
+        adjacency_.resize(static_cast<std::size_t>(number_of_nodes));
+        edges_.reserve(static_cast<std::size_t>(reserve_number_of_edges));
+    }
+
+    [[nodiscard]] NodeId number_of_nodes() const {
+        return number_of_nodes_;
+    }
+
+    [[nodiscard]] EdgeId number_of_edges() const {
+        return static_cast<EdgeId>(edges_.size());
+    }
+
+    [[nodiscard]] NodeId node_id_upper_bound() const {
+        return number_of_nodes_ == 0 ? 0 : number_of_nodes_ - 1;
+    }
+
+    [[nodiscard]] EdgeId edge_id_upper_bound() const {
+        return edges_.empty() ? 0 : static_cast<EdgeId>(edges_.size() - 1);
+    }
+
+    [[nodiscard]] std::vector<NodeId> nodes() const {
+        std::vector<NodeId> result(static_cast<std::size_t>(number_of_nodes_));
+        for (NodeId node = 0; node < number_of_nodes_; ++node) {
+            result[static_cast<std::size_t>(node)] = node;
+        }
+        return result;
+    }
+
+    [[nodiscard]] std::vector<EdgeId> edges() const {
+        std::vector<EdgeId> result(edges_.size());
+        for (EdgeId edge = 0; edge < edges_.size(); ++edge) {
+            result[static_cast<std::size_t>(edge)] = edge;
+        }
+        return result;
+    }
+
+    [[nodiscard]] Edge uv(const EdgeId edge) const {
+        validate_edge(edge);
+        return edges_[static_cast<std::size_t>(edge)];
+    }
+
+    [[nodiscard]] NodeId u(const EdgeId edge) const {
+        return uv(edge).first;
+    }
+
+    [[nodiscard]] NodeId v(const EdgeId edge) const {
+        return uv(edge).second;
+    }
+
+    [[nodiscard]] const std::vector<Edge> &uv_ids() const {
+        return edges_;
+    }
+
+    [[nodiscard]] const AdjacencyList &node_adjacency(const NodeId node) const {
+        validate_node(node);
+        return adjacency_[static_cast<std::size_t>(node)];
+    }
+
+    EdgeId insert_edge(const NodeId u, const NodeId v) {
+        validate_node(u);
+        validate_node(v);
+        if (u == v) {
+            throw std::invalid_argument("self edges are not supported");
+        }
+
+        const auto key = edge_key(u, v);
+        const auto found = edge_lookup_.find(key);
+        if (found != edge_lookup_.end()) {
+            return found->second;
+        }
+        return insert_new_edge(key.first, key.second);
+    }
+
+    [[nodiscard]] std::int64_t find_edge(const NodeId u, const NodeId v) const {
+        validate_node(u);
+        validate_node(v);
+        if (u == v) {
+            return -1;
+        }
+
+        const auto key = edge_key(u, v);
+        const auto found = edge_lookup_.find(key);
+        if (found == edge_lookup_.end()) {
+            return -1;
+        }
+        return static_cast<std::int64_t>(found->second);
+    }
+
+    [[nodiscard]] std::uint64_t serialization_size() const {
+        return 2 + 2 * number_of_edges();
+    }
+
+    [[nodiscard]] std::pair<std::vector<EdgeId>, std::vector<EdgeId>>
+    extract_subgraph_from_nodes(const std::vector<NodeId> &nodes) const {
+        std::unordered_set<NodeId> node_set;
+        node_set.reserve(nodes.size());
+        for (const auto node : nodes) {
+            validate_node(node);
+            node_set.insert(node);
+        }
+
+        std::vector<EdgeId> inner_edges;
+        std::vector<EdgeId> outer_edges;
+        for (const auto u : nodes) {
+            for (const auto adjacency : node_adjacency(u)) {
+                const auto v = adjacency.node;
+                const auto edge = adjacency.edge;
+                if (node_set.find(v) != node_set.end()) {
+                    if (u < v) {
+                        inner_edges.push_back(edge);
+                    }
+                } else {
+                    outer_edges.push_back(edge);
+                }
+            }
+        }
+        return {inner_edges, outer_edges};
+    }
+
+protected:
+    EdgeId insert_new_edge(const NodeId u, const NodeId v) {
+        const auto edge = static_cast<EdgeId>(edges_.size());
+        edges_.emplace_back(u, v);
+        edge_lookup_.emplace(Edge{u, v}, edge);
+        adjacency_[static_cast<std::size_t>(u)].push_back(Adjacency{v, edge});
+        adjacency_[static_cast<std::size_t>(v)].push_back(Adjacency{u, edge});
+        return edge;
+    }
+
+    void validate_node(const NodeId node) const {
+        if (node >= number_of_nodes_) {
+            throw std::out_of_range(
+                "node id must be < number_of_nodes, got node id=" +
+                std::to_string(node) + ", number_of_nodes=" +
+                std::to_string(number_of_nodes_)
+            );
+        }
+    }
+
+    void validate_edge(const EdgeId edge) const {
+        if (edge >= edges_.size()) {
+            throw std::out_of_range(
+                "edge id must be < number_of_edges, got edge id=" +
+                std::to_string(edge) + ", number_of_edges=" +
+                std::to_string(edges_.size())
+            );
+        }
+    }
+
+private:
+    struct EdgeHash {
+        std::size_t operator()(const Edge &edge) const {
+            const auto first = static_cast<std::size_t>(edge.first);
+            const auto second = static_cast<std::size_t>(edge.second);
+            return first ^ (second + 0x9e3779b97f4a7c15ULL + (first << 6U) + (first >> 2U));
+        }
+    };
+
+    static Edge edge_key(NodeId u, NodeId v) {
+        if (v < u) {
+            std::swap(u, v);
+        }
+        return {u, v};
+    }
+
+    NodeId number_of_nodes_;
+    std::vector<Edge> edges_;
+    std::vector<AdjacencyList> adjacency_;
+    std::unordered_map<Edge, EdgeId, EdgeHash> edge_lookup_;
+};
+
+} // namespace bioimage_cpp::graph