Store vertices in Graph

2 files changed, 12 insertions, 24 deletions

diff --git a/src/graph.h b/src/graph.h
index 8c04606..b4b300b 100644
--- a/src/graph.h
+++ b/src/graph.h
@@ -6,31 +6,30 @@
 
 class Graph {
 public:
-    Graph() : graph_() {}
+    Graph() : vertices_(), graph_() {}
 
     void add_vertex(int vertex) {
-        graph_[vertex] = std::vector<int>();
+        if (vertices_.find(vertex) == vertices_.end()) {
+            vertices_.insert(vertex);
+            graph_[vertex] = std::vector<int>();
+        }
     }
 
     void add_edge(int from, int to) {
+        add_vertex(from);
+        add_vertex(to);
         graph_[from].push_back(to);
     }
 
-    std::set<int> get_vertices() {
-        std::set<int> vertices;
-
-        for (auto vertex : graph_) {
-            vertices.insert(vertex.first);
-        }
-
-        return vertices;
+    const std::set<int> & get_vertices() const {
+        return vertices_;
     }
 
-
-    std::vector<int> const& get_neighbors(int vertex) {
-        return graph_[vertex];
+    const std::vector<int> & get_neighbors(int vertex) const {
+        return graph_.find(vertex)->second;
     }
 private:
+    std::set<int> vertices_;
     std::map<int, std::vector<int>> graph_;
 };
 
diff --git a/src/parse.h b/src/parse.h
index ec86302..52b8520 100644
--- a/src/parse.h
+++ b/src/parse.h
@@ -22,7 +22,6 @@ private:
     Graph graph_;
     std::ifstream input_file_;
     int number_edges_;
-    std::set<int> vertices_;
 
     void parse_number_edges_() {
         input_file_ >> number_edges_;
@@ -33,16 +32,6 @@ private:
 
         input_file_ >> from >> to;
 
-        if (vertices_.find(from) == vertices_.end()) {
-            graph_.add_vertex(from);
-            vertices_.insert(from);
-        }
-
-        if (vertices_.find(to) == vertices_.end()) {
-            graph_.add_vertex(to);
-            vertices_.insert(to);
-        }
-
         graph_.add_edge(from, to);
     }
 };