Add `Array#{permutation,combination}.

1 files changed, 96 insertions, 0 deletions

diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb
index e28e5238d..b525d3006 100644
--- a/mrbgems/mruby-array-ext/mrblib/array.rb
+++ b/mrbgems/mruby-array-ext/mrblib/array.rb
@@ -808,4 +808,100 @@ class Array
       n
     end
   end
+
+  ##
+  # call-seq:
+  #    ary.permutation { |p| block }          -> ary
+  #    ary.permutation                        -> Enumerator
+  #    ary.permutation(n) { |p| block }       -> ary
+  #    ary.permutation(n)                     -> Enumerator
+  #
+  # When invoked with a block, yield all permutations of length +n+ of the
+  # elements of the array, then return the array itself.
+  #
+  # If +n+ is not specified, yield all permutations of all elements.
+  #
+  # The implementation makes no guarantees about the order in which the
+  # permutations are yielded.
+  #
+  # If no block is given, an Enumerator is returned instead.
+  #
+  # Examples:
+  #
+  #  a = [1, 2, 3]
+  #  a.permutation.to_a    #=> [[1,2,3],[1,3,2],[2,1,3],[2,3,1],[3,1,2],[3,2,1]]
+  #  a.permutation(1).to_a #=> [[1],[2],[3]]
+  #  a.permutation(2).to_a #=> [[1,2],[1,3],[2,1],[2,3],[3,1],[3,2]]
+  #  a.permutation(3).to_a #=> [[1,2,3],[1,3,2],[2,1,3],[2,3,1],[3,1,2],[3,2,1]]
+  #  a.permutation(0).to_a #=> [[]] # one permutation of length 0
+  #  a.permutation(4).to_a #=> []   # no permutations of length 4
+  def permutation(n=self.size, &block)
+    size = self.size
+    return to_enum(:permutation, n) unless block
+    return if n > size
+    if n == 0
+       yield []
+    else
+      i = 0
+      while i<size
+        result = [self[i]]
+        if n-1 > 0
+          ary = self[0...i] + self[i+1..-1]
+          ary.permutation(n-1) do |c|
+            yield result + c
+          end
+        else
+          yield result
+        end
+        i += 1
+      end
+    end
+  end
+
+  ##
+  # call-seq:
+  #    ary.combination(n) { |c| block }    -> ary
+  #    ary.combination(n)                  -> Enumerator
+  #
+  # When invoked with a block, yields all combinations of length +n+ of elements
+  # from the array and then returns the array itself.
+  #
+  # The implementation makes no guarantees about the order in which the
+  # combinations are yielded.
+  #
+  # If no block is given, an Enumerator is returned instead.
+  #
+  # Examples:
+  #
+  #    a = [1, 2, 3, 4]
+  #    a.combination(1).to_a  #=> [[1],[2],[3],[4]]
+  #    a.combination(2).to_a  #=> [[1,2],[1,3],[1,4],[2,3],[2,4],[3,4]]
+  #    a.combination(3).to_a  #=> [[1,2,3],[1,2,4],[1,3,4],[2,3,4]]
+  #    a.combination(4).to_a  #=> [[1,2,3,4]]
+  #    a.combination(0).to_a  #=> [[]] # one combination of length 0
+  #    a.combination(5).to_a  #=> []   # no combinations of length 5
+
+  def combination(n, &block)
+    size = self.size
+    return to_enum(:combination, n) unless block
+    return if n > size
+    if n == 0
+       yield []
+    elsif n == 1
+      i = 0
+      while i<size
+        yield [self[i]]
+        i += 1
+      end
+    else
+      i = 0
+      while i<size
+        result = [self[i]]
+        self[i+1..-1].combination(n-1) do |c|
+          yield result + c
+        end
+        i += 1
+      end
+    end
+  end
 end