Adding weighted graph data structure

data_structures/graphs/weighted_graph.rb

@@ -0,0 +1,66 @@
+require 'set'
+
+##
+# This class aims to represent weighted graphs
+# (i.e. graphs for which edges between nodes have specific weights associated to them).
+#
+# Both directed (i.e. an edge between node U and node V does not imply an edge in the opposite direction)
+# and undirected graphs are supported, depending on the constructor invocation.
+
+class WeightedGraph
+  attr_reader :nodes
+  attr_reader :directed
+
+  def initialize(nodes: [], edges: {}, directed: true)
+    @nodes = Set[]
+    @edges = {}
+    @directed = directed
+    for node in nodes
+      add_node(node)
+    end
+    edges.each do |node, edges|
+      for neighbor, weight in edges
+        add_edge(node, neighbor, weight)
+      end
+    end
+  end
+
+  def add_node(node)
+    if include?(node)
+      raise ArgumentError, "node #{node} already exists in this graph!"
+    end
+    @nodes.add(node)
+    @edges[node] = {}
+  end
+
+  def add_edge(start_node, end_node, weight)
+    if has_neighbor?(start_node, end_node)
+      raise ArgumentError, "node #{start_node} already has an edge to #{end_node} in this graph!"
+    end
+    @edges[start_node][end_node] = weight
+    @edges[end_node][start_node] = weight unless directed
+  end
+
+  def edges(node)
+    unless include?(node)
+      raise ArgumentError, "node #{node} does not exist in this graph!"
+    end
+    @edges[node]
+  end
+
+  def empty?
+    nodes.empty?
+  end
+
+  def include?(node)
+    nodes.include?(node)
+  end
+
+  def has_neighbor?(start_node, end_node)
+    edges(start_node).key?(end_node)
+  end
+
+  def edge_weight(start_node, end_node)
+    edges(start_node)[end_node]
+  end
+end

data_structures/graphs/weighted_graph_test.rb

@@ -0,0 +1,88 @@
+require 'minitest/autorun'
+require 'set'
+require_relative 'weighted_graph'
+
+class TestWeightedGraph < Minitest::Test
+  def test_directed_weighted_graph_creation
+    graph = WeightedGraph.new(nodes: [:u, :v, :w], edges: {:u => [[:v, 1]]}, directed: true)
+
+    assert graph.nodes.to_set == Set[:u, :v, :w]
+    assert graph.edges(:u) == {:v => 1}
+    assert graph.edges(:v).empty?
+    assert graph.edges(:w).empty?
+  end
+
+  def test_undirected_weighted_graph_creation
+    graph = WeightedGraph.new(nodes: [:u, :v, :w], edges: {:u => [[:v, 1]]}, directed: false)
+
+    assert graph.nodes.to_set == Set[:u, :v, :w]
+    assert graph.edges(:u) == {:v => 1}
+    assert graph.edges(:v) == {:u => 1}
+    assert graph.edges(:w).empty?
+  end
+
+  def test_empty_returns_true_for_empty_graph
+    graph = WeightedGraph.new
+
+    assert graph.empty?
+  end
+
+  def test_empty_returns_false_for_non_empty_graph
+    graph = WeightedGraph.new(nodes: [:u])
+
+    assert !graph.empty?
+  end
+
+  def test_include_returns_true_for_graph_nodes
+    graph = WeightedGraph.new(nodes: [:u])
+
+    assert graph.include?(:u)
+  end
+
+  def test_include_returns_false_for_non_graph_nodes
+    graph = WeightedGraph.new
+
+    assert !graph.include?(:u)
+  end
+
+  def test_has_neighbor_returns_true_for_graph_node_neighbors
+    graph = WeightedGraph.new(nodes: [:u, :v], edges: {:u => [[:v, 1]]})
+
+    assert graph.has_neighbor?(:u, :v)
+  end
+
+  def test_has_neighbor_returns_false_for_non_graph_node_neighbors
+    graph = WeightedGraph.new(nodes: [:u, :v])
+
+    assert !graph.has_neighbor?(:u, :v)
+  end
+
+  def test_edge_weight_returns_neighbor_edge_weight
+    graph = WeightedGraph.new(nodes: [:u, :v], edges: {:u => [[:v, 4]]})
+
+    assert graph.edge_weight(:u, :v) == 4
+  end
+
+  def test_add_node_adds_node_to_graph
+    graph = WeightedGraph.new
+    graph.add_node(:u)
+
+    assert graph.nodes.to_set == Set[:u]
+  end
+
+  def test_add_edge_adds_edge_to_directed_weighted_graph
+    graph = WeightedGraph.new(nodes: [:u, :v], directed: true)
+    graph.add_edge(:u, :v, 2)
+
+    assert graph.edges(:u) == {:v => 2}
+    assert graph.edges(:v).empty?
+  end
+
+  def test_add_edge_adds_edge_to_directed_weighted_graph
+    graph = WeightedGraph.new(nodes: [:u, :v], directed: false)
+    graph.add_edge(:u, :v, 2)
+
+    assert graph.edges(:u) == {:v => 2}
+    assert graph.edges(:v) == {:u => 2}
+  end
+end