Add vertex weights and neighbor iterator for WeightedEdgeGraph

include/aoc/common/WeightedEdgeGraph.hpp

@@ -15,6 +15,7 @@
 
 #pragma once
 
+#include <iterator>
 #include <vector>
 
 #include <aoc/common/VertexEdgeIncidence.hpp>
@@ -23,11 +24,90 @@ class WeightedEdgeGraph
 {
 public:
     WeightedEdgeGraph();
-    int addVertex();
+    int addVertex(const int weight);
     void addEdge(const int vertex1, const int vertex2, const int weight);
+    int getVertexWeight(const int vertex) const;
+    int getEdgeWeight(const int edge) const;
     int dijkstra(const int source, const int target) const;
+
+    struct Incidence
+    {
+    public:
+        Incidence(const int vertex, const int edge)
+            : vertex{ vertex }, edge{ edge }
+        {
+        }
+        int vertex;
+        int edge;
+    };
+
+    struct NeighborIterator
+    {
+    public:
+        // Iterator traits.
+        using difference_type = std::ptrdiff_t;
+        using value_type = Incidence;
+        using pointer = const value_type*;
+        using reference = const value_type&;
+        using iterator_category = std::forward_iterator_tag;
+
+        NeighborIterator(const WeightedEdgeGraph& graph, const int vertex)
+            : graph_{ graph }, incidence_{ -1 }, value_{ -1, -1 }
+        {
+            if (vertex >= 0 && vertex < graph_.firstVertexIncidences_.size())
+            {
+                setIncidence(graph_.firstVertexIncidences_[vertex]);
+          }
+        };
+        NeighborIterator& operator++()
+        {
+            if (incidence_ >= 0)
+            {
+                setIncidence(graph_.vertexEdgeIncidences_[incidence_].next);
+            }
+            return *this;
+        }
+        NeighborIterator operator++(int)
+        {
+            NeighborIterator it = *this;
+            ++(*this);
+            return it;
+        }
+        bool operator==(NeighborIterator other) const
+        {
+            return incidence_ == other.incidence_;
+        }
+        bool operator!=(NeighborIterator other) const
+        {
+            return !(*this == other);
+        }
+        const reference operator*() const
+        {
+            return value_;
+        }
+        const pointer operator->() const
+        {
+            return &value_;
+        }
+    private:
+        const WeightedEdgeGraph& graph_;
+        int incidence_;
+        value_type value_;
+        void setIncidence(const int incidence)
+        {
+            incidence_ = incidence;
+            if (incidence_ >= 0)
+            {
+                value_ = { graph_.vertexEdgeIncidences_[incidence_].vertex, incidence_ >> 1 };
+            }
+        }
+    };
+
+    NeighborIterator begin(const int vertex) const;
+    NeighborIterator end() const;
 private:
     std::vector<int> firstVertexIncidences_;
     std::vector<VertexEdgeIncidence> vertexEdgeIncidences_;
     std::vector<int> edgeWeights_;
+    std::vector<int> vertexWeights_;
 };

src/common/WeightedEdgeGraph.cpp

@@ -23,9 +23,10 @@ WeightedEdgeGraph::WeightedEdgeGraph()
 {
 }
 
-int WeightedEdgeGraph::addVertex()
+int WeightedEdgeGraph::addVertex(const int weight)
 {
     firstVertexIncidences_.push_back(-1);
+    vertexWeights_.push_back(weight);
     return (int)firstVertexIncidences_.size() - 1;
 }
 
@@ -38,6 +39,16 @@ void WeightedEdgeGraph::addEdge(const int vertex1, const int vertex2, const int
     edgeWeights_.push_back(weight);
 }
 
+int WeightedEdgeGraph::getVertexWeight(const int vertex) const
+{
+    return vertexWeights_[vertex];
+}
+
+int WeightedEdgeGraph::getEdgeWeight(const int edge) const
+{
+    return edgeWeights_[edge];
+}
+
 int WeightedEdgeGraph::dijkstra(const int source, const int target) const
 {
     std::vector<int> distances(firstVertexIncidences_.size(), std::numeric_limits<int>::max());
@@ -72,3 +83,13 @@ int WeightedEdgeGraph::dijkstra(const int source, const int target) const
 
     return -1;
 }
+
+WeightedEdgeGraph::NeighborIterator WeightedEdgeGraph::begin(const int vertex) const
+{
+    return { *this, vertex };
+}
+
+WeightedEdgeGraph::NeighborIterator WeightedEdgeGraph::end() const
+{
+    return { *this, -1 };
+}