diff --git a/lib/networkx/graph.rb b/lib/networkx/graph.rb index 84d5be9..25a632c 100644 --- a/lib/networkx/graph.rb +++ b/lib/networkx/graph.rb @@ -1,40 +1,44 @@ module NetworkX class Graph def initialize - @nodes = Set.new - @edges = {} + @adj = {} + end + + def add_node(node) + @adj[node] = {} unless @adj.key? node + end + + def add_nodes_from(nodes) + nodes.each { |node| add_node(node) } end def add_edge(node_1, node_2, **attrs) add_node(node_1) add_node(node_2) - edge_key = Set.new([node_1, node_2]) - @edges[edge_key] = attrs + @adj[node_1][node_2] = attrs end - def add_edges(edges) - case edges - when Array - edges.each { |node_1, node_2, **attrs| add_edge(node_1, node_2, **attrs) } - else - raise ArgumentError, 'Expected argument to be an Array of edges, '\ - "received #{edges.class.name} instead." - end + def add_edges_from(edges) + edges.each { |edge| add_edge(*edge) } end - def add_node(node) - @nodes.add(node) + def remove_node(node) + @adj.delete(node) + @adj.each_key { |key| @adj[key].delete(node) } + end + + def remove_nodes_from(nodes) + nodes.each { |node| @adj.delete(node) } + nodes.each { |node| @adj.each_key { |key| @adj[key].delete(node) } } + end + + def remove_edge(node_1, node_2) + @adj[node_1].delete(node_2) if @adj.key? node_1 end - def add_nodes(nodes) - case nodes - when Set, Array - nodes.each { |node| add_node(node) } - else - raise ArgumentError, 'Expected argument to be an Array or Set of nodes, '\ - "received #{nodes.class.name} instead." - end + def remove_edges_from(edges) + edges.each { |edge| remove_edge(*edge) } end end end