Merge c35e341 into e6a24a0

gabynaiman · Oct 28, 2020 · 775407b · 775407b
2 parents e6a24a0 + c35e341
commit 775407b
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/lib/hierarchical_graph.rb b/lib/hierarchical_graph.rb
@@ -32,8 +32,8 @@ def roots
 
   def add_node(id, attributes={})
     clear_cache
-    parent_to_children[id] = Set.new
-    child_to_parents[id] = Set.new
+    parent_to_children[id] ||= Set.new
+    child_to_parents[id] ||= Set.new
     nodes[id] = Node.new self, id, attributes
   end
 
@@ -73,11 +73,11 @@ def remove_relation(parent_id:, child_id:)
   end
 
   def parents_of(id)
-    child_to_parents[id].map { |node_id| nodes[node_id] }
+    child_to_parents.fetch(id, Set.new).map { |node_id| nodes[node_id] }
   end
 
   def children_of(id)
-    parent_to_children[id].map { |node_id| nodes[node_id] }
+    parent_to_children.fetch(id, Set.new).map { |node_id| nodes[node_id] }
   end
 
   def ancestors_of(id)