lib/graph.rb in graph-rb-0.1.0 vs lib/graph.rb in graph-rb-0.1.1

@@ -13,10 +13,15 @@
   def include?(vertex_or_edge)
     return include_vertex?(vertex_or_edge) if vertex_or_edge.is_a?(Graph::Vertex)
     return include_edge?(vertex_or_edge) if vertex_or_edge.is_a?(Graph::Edge)
   end
 
+  def index_of(vertex_or_edge)
+    return index_of_vertex(vertex_or_edge) if vertex_or_edge.is_a?(Graph::Vertex)
+    return index_of_edge(vertex_or_edge) if vertex_or_edge.is_a?(Graph::Edge)
+  end
+
   def add(vertex_or_edge)
     return add_vertex(vertex_or_edge) if vertex_or_edge.is_a?(Graph::Vertex)
     return add_edge(vertex_or_edge) if vertex_or_edge.is_a?(Graph::Edge)
   end
 
@@ -24,19 +29,18 @@
     return delete_vertex(vertex_or_edge) if vertex_or_edge.is_a?(Graph::Vertex)
     return delete_edge(vertex_or_edge) if vertex_or_edge.is_a?(Graph::Edge)
   end
 
   def to_h
-    result = {
-      vertices: [],
-      edges: []
-    }
+    # O(vertices.size) + O(edges.size)
+    result = {}
     @vertices.each do |vertex|
-      result[:vertices] << vertex.uid
+      result[vertex.key] = vertex.to_h
     end
     @edges.each do |edge|
-      result[:edges] << { from: edge.from.uid, to: edge.to.uid }
+      result[edge.from.key][:edges] << edge.to_h
+      result[edge.to.key][:edges] << edge.to_h
     end
     result
   end
 
   private
@@ -47,41 +51,54 @@
 
   def include_edge?(edge)
     @edges.include? edge
   end
 
+  def index_of_vertex(target)
+    @vertices.index { |vertex| vertex.key == target.key }
+  end
+
+  def index_of_edge(target)
+    @edges.index { |edge| edge.key == target.key }
+  end
+
   def add_vertex(vertex)
-    return false if include_vertex?(vertex)
+    return false if index_of_vertex(vertex)
 
     vertex.send(:add_to, self)
     @vertices << vertex
     self
   end
 
   def add_edge(edge)
-    return false if include_edge?(edge)
+    return false if index_of_edge(edge)
 
-    edge.send(:add_to, self)
     add_vertex(edge.from)
     add_vertex(edge.to)
+    edge.send(:add_to, self)
     @edges << edge
     self
   end
 
   def delete_vertex(vertex)
-    return false unless include_vertex?(vertex)
+    idx = index_of_vertex vertex
+    return false unless idx
+
     edges.each do |edge|
-      delete_edge(edge) if (edge.from == vertex) || (edge.to == vertex)
+      delete_edge(edge) if (edge.from.key == vertex.key) || (edge.to.key == vertex.key)
     end
     vertex.send(:add_to, nil)
-    @vertices.delete(vertex)
+    @vertices.delete_at idx
     self
   end
 
   def delete_edge(edge)
-    return false unless include_edge?(edge)
-    
+    idx = index_of_edge edge
+    return false unless idx
+
     edge.send(:add_to, nil)
-    @edges.delete(edge)
+    @edges.delete_at idx
     self
   end
-end
\ No newline at end of file
+
+  alias << add
+end