* lib/matrix.rb: Matrix#each{_with_index} can iterate over a subset of the Matrix

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@32277 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

1 files changed, 110 insertions, 19 deletions

diff --git a/lib/matrix.rb b/lib/matrix.rb
index 3fda08b18d..6c037063d3 100644
--- a/lib/matrix.rb
+++ b/lib/matrix.rb
@@ -365,36 +365,127 @@ class Matrix
 
   #
   # Yields all elements of the matrix, starting with those of the first row,
-  # or returns an Enumerator is no block given
+  # or returns an Enumerator is no block given.
+  # Elements can be restricted by passing an argument:
+  # * :all (default): yields all elements
+  # * :diagonal: yields only elements on the diagonal
+  # * :off_diagonal: yields all elements except on the diagonal
+  # * :lower: yields only elements on or below the diagonal
+  # * :strict_lower: yields only elements below the diagonal
+  # * :strict_upper: yields only elements above the diagonal
+  # * :upper: yields only elements on or above the diagonal
+  #
   #   Matrix[ [1,2], [3,4] ].each { |e| puts e }
   #     # => prints the numbers 1 to 4
-  #
-  def each(&block) # :yield: e
-    return to_enum(:each) unless block_given?
-    @rows.each do |row|
-      row.each(&block)
+  #   Matrix[ [1,2], [3,4] ].each(:strict_lower).to_a # => [3]
+  #
+  def each(which = :all) # :yield: e
+    return to_enum :each, which unless block_given?
+    last = column_size - 1
+    case which
+    when :all
+      block = Proc.new
+      @rows.each do |row|
+        row.each(&block)
+      end
+    when :diagonal
+      @rows.each_with_index do |row, row_index|
+        yield row.fetch(row_index){return self}
+      end
+    when :off_diagonal
+      @rows.each_with_index do |row, row_index|
+        column_size.times do |col_index|
+          yield row[col_index] unless row_index == col_index
+        end
+      end
+    when :lower
+      @rows.each_with_index do |row, row_index|
+        0.upto([row_index, last].min) do |col_index|
+          yield row[col_index]
+        end
+      end
+    when :strict_lower
+      @rows.each_with_index do |row, row_index|
+        [row_index, column_size].min.times do |col_index|
+          yield row[col_index]
+        end
+      end
+    when :strict_upper
+      @rows.each_with_index do |row, row_index|
+        (row_index+1).upto(last) do |col_index|
+          yield row[col_index]
+        end
+      end
+    when :upper
+      @rows.each_with_index do |row, row_index|
+        row_index.upto(last) do |col_index|
+          yield row[col_index]
+        end
+      end
+    else
+      Matrix.Raise ArgumentError, "expected #{which.inspect} to be one of :all, :diagonal, :off_diagonal, :lower, :strict_lower, :strict_upper or :upper"
     end
     self
   end
 
   #
-  # Yields all elements of the matrix, starting with those of the first row,
-  # along with the row index and column index,
-  # or returns an Enumerator is no block given
+  # Same as #each, but the row index and column index in addition to the element
+  #
   #   Matrix[ [1,2], [3,4] ].each_with_index do |e, row, col|
   #     puts "#{e} at #{row}, #{col}"
   #   end
-  #     # => 1 at 0, 0
-  #     # => 2 at 0, 1
-  #     # => 3 at 1, 0
-  #     # => 4 at 1, 1
-  #
-  def each_with_index(&block) # :yield: e, row, column
-    return to_enum(:each_with_index) unless block_given?
-    @rows.each_with_index do |row, row_index|
-      row.each_with_index do |e, col_index|
-        yield e, row_index, col_index
+  #     # => Prints:
+  #     #    1 at 0, 0
+  #     #    2 at 0, 1
+  #     #    3 at 1, 0
+  #     #    4 at 1, 1
+  #
+  def each_with_index(which = :all) # :yield: e, row, column
+    return to_enum :each_with_index, which unless block_given?
+    last = column_size - 1
+    case which
+    when :all
+      @rows.each_with_index do |row, row_index|
+        row.each_with_index do |e, col_index|
+          yield e, row_index, col_index
+        end
+      end
+    when :diagonal
+      @rows.each_with_index do |row, row_index|
+        yield row.fetch(row_index){return self}, row_index, row_index
+      end
+    when :off_diagonal
+      @rows.each_with_index do |row, row_index|
+        column_size.times do |col_index|
+          yield row[col_index], row_index, col_index unless row_index == col_index
+        end
+      end
+    when :lower
+      @rows.each_with_index do |row, row_index|
+        0.upto([row_index, last].min) do |col_index|
+          yield row[col_index], row_index, col_index
+        end
       end
+    when :strict_lower
+      @rows.each_with_index do |row, row_index|
+        [row_index, column_size].min.times do |col_index|
+          yield row[col_index], row_index, col_index
+        end
+      end
+    when :strict_upper
+      @rows.each_with_index do |row, row_index|
+        (row_index+1).upto(last) do |col_index|
+          yield row[col_index], row_index, col_index
+        end
+      end
+    when :upper
+      @rows.each_with_index do |row, row_index|
+        row_index.upto(last) do |col_index|
+          yield row[col_index], row_index, col_index
+        end
+      end
+    else
+      Matrix.Raise ArgumentError, "expected #{which.inspect} to be one of :all, :diagonal, :off_diagonal, :lower, :strict_lower, :strict_upper or :upper"
     end
     self
   end