Merge pull request #5598 from dearblue/array-combinations

Add `Array#{repeated_combination,repeated_permutation}` methods

doc/mruby3.1.md

@@ -37,6 +37,8 @@ Keyword arguments are basically separated from ordinal arguments
 ## New Methods
 
 * `Array#product`
+* `Array#repeated_combination`
+* `Array#repeated_permutation`
 * `Random.bytes`
 * `Random#bytes`
 

mrbgems/mruby-array-ext/mrblib/array.rb

@@ -943,4 +943,77 @@ class Array
 
     result
   end
+
+  ##
+  # call-seq:
+  #   ary.repeated_combination(n) { |combination| ... }   ->   self
+  #   ary.repeated_combination(n)                         ->   enumerator
+  #
+  # A +combination+ method that contains the same elements.
+  def repeated_combination(n, &block)
+    raise TypeError, "no implicit conversion into Integer" unless 0 <=> n
+    return to_enum(:repeated_combination, n) unless block
+    __repeated_combination(n, false, &block)
+  end
+
+  ##
+  # call-seq:
+  #   ary.repeated_permutation(n) { |permutation| ... }   ->   self
+  #   ary.repeated_permutation(n)                         ->   enumerator
+  #
+  # A +permutation+ method that contains the same elements.
+  def repeated_permutation(n, &block)
+    raise TypeError, "no implicit conversion into Integer" unless 0 <=> n
+    return to_enum(:repeated_permutation, n) unless block
+    __repeated_combination(n, true, &block)
+  end
+
+  def __repeated_combination(n, permutation, &block)
+    case n
+    when 0
+      yield []
+    when 1
+      i = 0
+      while i < self.size
+        yield [self[i]]
+        i += 1
+      end
+    else
+      if n > 0
+        v = [0] * n
+        while true
+          tmp = [nil] * n
+          i = 0
+          while i < n
+            tmp[i] = self[v[i]]
+            i += 1
+          end
+
+          yield tmp
+
+          tmp = self.size
+          i = n - 1
+          while i >= 0
+            v[i] += 1
+            break if v[i] < tmp
+            i -= 1
+          end
+          break unless v[0] < tmp
+          i = 1
+          while i < n
+            unless v[i] < tmp
+              if permutation
+                v[i] = 0
+              else
+                v[i] = v[i - 1]
+              end
+            end
+            i += 1
+          end
+        end
+      end
+    end
+
+    self
+  end
 end

mrbgems/mruby-array-ext/test/array.rb

@@ -18,6 +18,14 @@ def assert_combination(exp, receiver, *args)
   assert_permutation_combination(exp, receiver, :combination, *args)
 end
 
+def assert_repeated_permutation(exp, receiver, *args)
+  assert_permutation_combination(exp, receiver, :repeated_permutation, *args)
+end
+
+def assert_repeated_combination(exp, receiver, *args)
+  assert_permutation_combination(exp, receiver, :repeated_combination, *args)
+end
+
 assert("Array#assoc") do
   s1 = [ "colors", "red", "blue", "green" ]
   s2 = [ "letters", "a", "b", "c" ]
@@ -443,3 +451,51 @@ assert "Array#product" do
   assert_equal base, base.product([4, 5, 6], [7, 8, 9]) { |e| x << e }
   assert_equal expect, x
 end
+
+assert("Array#repeated_combination") do
+  a = [1, 2, 3]
+  assert_raise(ArgumentError) { a.repeated_combination }
+  #assert_kind_of(Enumerator, a.repeated_combination(1))
+  assert_repeated_combination([[1],[2],[3]], a, 1)
+  assert_repeated_combination([[1,1],[1,2],[1,3],[2,2],[2,3],[3,3]], a, 2)
+  assert_repeated_combination([[1,1,1],[1,1,2],[1,1,3],[1,2,2],[1,2,3],[1,3,3],[2,2,2],
+                               [2,2,3],[2,3,3],[3,3,3]], a, 3)
+  assert_repeated_combination([[1,1,1,1],[1,1,1,2],[1,1,1,3],[1,1,2,2],[1,1,2,3],[1,1,3,3],
+                               [1,2,2,2],[1,2,2,3],[1,2,3,3],[1,3,3,3],[2,2,2,2],[2,2,2,3],
+                               [2,2,3,3],[2,3,3,3],[3,3,3,3]], a, 4)
+  assert_repeated_combination([[1,1,1,1,1],[1,1,1,1,2],[1,1,1,1,3],[1,1,1,2,2],[1,1,1,2,3],
+                               [1,1,1,3,3],[1,1,2,2,2],[1,1,2,2,3],[1,1,2,3,3],[1,1,3,3,3],
+                               [1,2,2,2,2],[1,2,2,2,3],[1,2,2,3,3],[1,2,3,3,3],[1,3,3,3,3],
+                               [2,2,2,2,2],[2,2,2,2,3],[2,2,2,3,3],[2,2,3,3,3],[2,3,3,3,3],
+                               [3,3,3,3,3]], a, 5)
+  assert_repeated_combination([[]], a, 0)
+  assert_repeated_combination([], a, -1)
+end
+
+assert("Array#repeated_permutation") do
+  a = [1, 2, 3]
+  assert_raise(ArgumentError) { a.repeated_permutation }
+  #assert_kind_of(Enumerator, a.repeated_permutation(1))
+  assert_repeated_permutation([[1],[2],[3]], a, 1)
+  assert_repeated_permutation([[1,1],[1,2],[1,3],[2,1],[2,2],[2,3],[3,1],[3,2],[3,3]], a, 2)
+  assert_repeated_permutation([[1,1,1],[1,1,2],[1,1,3],[1,2,1],[1,2,2],[1,2,3],[1,3,1],[1,3,2],[1,3,3],
+                               [2,1,1],[2,1,2],[2,1,3],[2,2,1],[2,2,2],[2,2,3],[2,3,1],[2,3,2],[2,3,3],
+                               [3,1,1],[3,1,2],[3,1,3],[3,2,1],[3,2,2],[3,2,3],[3,3,1],[3,3,2],[3,3,3]],
+                              a, 3)
+  assert_repeated_permutation([[1,1,1,1],[1,1,1,2],[1,1,1,3],[1,1,2,1],[1,1,2,2],[1,1,2,3],
+                               [1,1,3,1],[1,1,3,2],[1,1,3,3],[1,2,1,1],[1,2,1,2],[1,2,1,3],
+                               [1,2,2,1],[1,2,2,2],[1,2,2,3],[1,2,3,1],[1,2,3,2],[1,2,3,3],
+                               [1,3,1,1],[1,3,1,2],[1,3,1,3],[1,3,2,1],[1,3,2,2],[1,3,2,3],
+                               [1,3,3,1],[1,3,3,2],[1,3,3,3],[2,1,1,1],[2,1,1,2],[2,1,1,3],
+                               [2,1,2,1],[2,1,2,2],[2,1,2,3],[2,1,3,1],[2,1,3,2],[2,1,3,3],
+                               [2,2,1,1],[2,2,1,2],[2,2,1,3],[2,2,2,1],[2,2,2,2],[2,2,2,3],
+                               [2,2,3,1],[2,2,3,2],[2,2,3,3],[2,3,1,1],[2,3,1,2],[2,3,1,3],
+                               [2,3,2,1],[2,3,2,2],[2,3,2,3],[2,3,3,1],[2,3,3,2],[2,3,3,3],
+                               [3,1,1,1],[3,1,1,2],[3,1,1,3],[3,1,2,1],[3,1,2,2],[3,1,2,3],
+                               [3,1,3,1],[3,1,3,2],[3,1,3,3],[3,2,1,1],[3,2,1,2],[3,2,1,3],
+                               [3,2,2,1],[3,2,2,2],[3,2,2,3],[3,2,3,1],[3,2,3,2],[3,2,3,3],
+                               [3,3,1,1],[3,3,1,2],[3,3,1,3],[3,3,2,1],[3,3,2,2],[3,3,2,3],
+                               [3,3,3,1],[3,3,3,2],[3,3,3,3]], a, 4)
+  assert_repeated_permutation([[]], a, 0)
+  assert_repeated_permutation([], a, -1)
+end