diff --git a/day_04/2.rb b/day_04/2.rb
index ec39356..5c79a6d 100755
--- a/day_04/2.rb
+++ b/day_04/2.rb
@@ -12,24 +12,25 @@ class Main
 
     @sum = 0
 
-    removed_sum = 1
     grid = Grid.new(data)
-    until removed_sum == 0
-      removed_sum = 0
+    loop do
+      grid_updated = false
       (0...grid.rows.size).each do |row_index|
         (0...grid.columns.size).each do |column_index|
-          next unless grid.cell_at(row_index, column_index).value == "@"
+          cell = grid.cell_at(row_index, column_index)
+          next unless cell.value == "@"
 
-          memo = grid.cell_at(row_index, column_index).neighbours.reject do |cell|
-            grid.cell_at(cell.row, cell.column).value != "@"
-          end
-          if memo.size < 4
+          if cell.neighbours.count do |neighbour|
+            neighbour.value == "@"
+          end < 4
             data[row_index][column_index] = "x"
-            removed_sum += 1
+            @sum += 1
+            grid_updated = true
           end
         end
       end
-      @sum += removed_sum
+
+      break unless grid_updated
       grid = Grid.new(data)
     end