app/services/tables/distribution.rb

@@ -13,7 +13,7 @@ module Tables
     def random_distribution(people)
       @tables = []
 
-      @tables << people.slice!(0..rand(@min_per_table..@max_per_table)) while people.any?
+      @tables << Table.new(people.slice!(0..rand(@min_per_table..@max_per_table))) while people.any?
     end
 
     def discomfort
@@ -59,7 +59,7 @@ module Tables
     private
 
     def local_discomfort(table)
-      DiscomfortCalculator.new(table).calculate
+      table.discomfort ||= DiscomfortCalculator.new(table).calculate
     end
   end
 end

app/services/tables/swap.rb

@@ -8,6 +8,9 @@ module Tables
     def each
       @initial_solution.tables.combination(2) do |table_a, table_b|
         table_a.product(table_b).each do |(person_a, person_b)|
+          original_discomfort_a = table_a.reset
+          original_discomfort_b = table_b.reset
+
           table_a.delete(person_a)
           table_b.delete(person_b)
 
@@ -21,6 +24,9 @@ module Tables
 
           table_a << person_a
           table_b << person_b
+
+          table_a.discomfort = original_discomfort_a
+          table_b.discomfort = original_discomfort_b
         end
       end
     end

app/services/tables/table.rb

@@ -0,0 +1,15 @@
+module Tables
+  class Table < Array
+    attr_accessor :discomfort
+    def initialize(*args)
+      super
+      reset
+    end
+
+    def reset
+      original_discomfort = discomfort
+      @discomfort = nil
+      original_discomfort
+    end
+  end
+end

config/initializers/ruby_extensions.rb

@@ -3,3 +3,9 @@ class Numeric
     Money.from_amount(self, "EUR").format
   end
 end
+
+class Array
+  def to_table
+    Tables::Table.new(self)
+  end
+end

spec/services/tables/swap_spec.rb

@@ -14,8 +14,8 @@ module Tables
       context 'when there are two tables with two people each' do
         let(:initial_solution) do
           Distribution.new(min_per_table: 2, max_per_table: 2).tap do |distribution|
-            distribution.tables << %i[a b]
+            distribution.tables << %i[a b].to_table
-            distribution.tables << %i[c d]
+            distribution.tables << %i[c d].to_table
           end
         end
 
@@ -32,8 +32,8 @@ module Tables
       context 'when there are two tables with three people each' do
         let(:initial_solution) do
           Distribution.new(min_per_table: 3, max_per_table: 3).tap do |distribution|
-            distribution.tables << %i[a b c]
+            distribution.tables << %i[a b c].to_table
-            distribution.tables << %i[d e f]
+            distribution.tables << %i[d e f].to_table
           end
         end
 
@@ -55,9 +55,9 @@ module Tables
       context 'when there are three tables with two people each' do
         let(:initial_solution) do
           Distribution.new(min_per_table: 2, max_per_table: 2).tap do |distribution|
-            distribution.tables << %i[a b]
+            distribution.tables << %i[a b].to_table
-            distribution.tables << %i[c d]
+            distribution.tables << %i[c d].to_table
-            distribution.tables << %i[e f]
+            distribution.tables << %i[e f].to_table
           end
         end