aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
-rw-r--r--array.c13
-rw-r--r--test/ruby/test_enumerator.rb5
2 files changed, 17 insertions, 1 deletions
diff --git a/array.c b/array.c
index 3eeafc8621..03d26afdec 100644
--- a/array.c
+++ b/array.c
@@ -4645,6 +4645,17 @@ rcombinate0(long n, long r, long *p, long index, long rest, VALUE values)
}
}
+static VALUE
+rb_ary_repeated_combination_size(VALUE ary, VALUE args)
+{
+ long n = RARRAY_LEN(ary);
+ long k = NUM2LONG(RARRAY_PTR(args)[0]);
+ if (k == 0) {
+ return LONG2FIX(1);
+ }
+ return binomial_coefficient(k, n + k - 1);
+}
+
/*
* call-seq:
* ary.repeated_combination(n) { |c| block } -> ary
@@ -4678,7 +4689,7 @@ rb_ary_repeated_combination(VALUE ary, VALUE num)
long n, i, len;
n = NUM2LONG(num); /* Combination size from argument */
- RETURN_ENUMERATOR(ary, 1, &num); /* Return Enumerator if no block */
+ RETURN_SIZED_ENUMERATOR(ary, 1, &num, rb_ary_repeated_combination_size); /* Return enumerator if no block */
len = RARRAY_LEN(ary);
if (n < 0) {
/* yield nothing */
diff --git a/test/ruby/test_enumerator.rb b/test/ruby/test_enumerator.rb
index 1c04a05bb0..93d46a6c23 100644
--- a/test/ruby/test_enumerator.rb
+++ b/test/ruby/test_enumerator.rb
@@ -450,6 +450,11 @@ class TestEnumerator < Test::Unit::TestCase
check_consistency_for_combinatorics(:repeated_permutation)
assert_equal 291733167875766667063796853374976,
(1..42).to_a.repeated_permutation(20).size # 42 ** 20
+
+ check_consistency_for_combinatorics(:repeated_combination)
+ assert_equal 28258808871162574166368460400,
+ (1..59).to_a.repeated_combination(42).size
+ # 1.upto(100).inject(:*) / 1.upto(42).inject(:*) / 1.upto(58).inject(:*)
end
end