Merge pull request #6817 from dearblue/array-ext

Integrate `Array#{permutation,combination}` into `Array#__combination`
This commit is contained in:
Yukihiro "Matz" Matsumoto
2026-05-08 10:56:42 +09:00
committed by GitHub
2 changed files with 110 additions and 76 deletions
+8 -52
View File
@@ -450,27 +450,7 @@ class Array
# 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)
n = n.__to_int
return to_enum(:permutation, n) unless block
size = self.size
if n == 0
yield []
elsif 0 < n && n <= size
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
self
__combination(:permutation, n, &block)
end
##
@@ -497,28 +477,7 @@ class Array
# a.combination(5).to_a #=> [] # no combinations of length 5
def combination(n, &block)
n = n.__to_int
return to_enum(:combination, n) unless block
size = self.size
if n == 0
yield []
elsif n == 1
i = 0
while i<size
yield [self[i]]
i += 1
end
elsif n <= size
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
self
__combination(:combination, n, &block)
end
##
@@ -642,9 +601,7 @@ class Array
# a = [1, 2, 3]
# a.repeated_combination(2).to_a #=> [[1,1],[1,2],[1,3],[2,2],[2,3],[3,3]]
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)
__combination(:repeated_combination, n, &block)
end
##
@@ -667,14 +624,13 @@ class Array
# a = [1, 2]
# a.repeated_permutation(2).to_a #=> [[1,1],[1,2],[2,1],[2,2]]
def repeated_permutation(n, &block)
n = n.__to_int
raise TypeError, "no implicit conversion into Integer" unless 0 <=> n
return to_enum(:repeated_permutation, n) unless block
__repeated_combination(n, true, &block)
__combination(:repeated_permutation, n, &block)
end
def __repeated_combination(k, permutation, &block)
def __combination(mode, k, &block)
k = k.__to_int
return to_enum(mode, k) unless block
case k
when 0
yield []
@@ -686,7 +642,7 @@ class Array
i += 1
end
else
if state = __combination_init(k, permutation)
if state = __combination_init(mode, k)
# Use C iterator for complex cases
while tmp = __combination_next(state)
yield tmp
+102 -24
View File
@@ -33,8 +33,7 @@ typedef khash_t(ary_set) ary_set_t;
struct mrb_combination_state {
mrb_int *indices;
mrb_int n, k; /* nPk, nCk */
mrb_bool permutation;
mrb_bool finished;
int mode;
};
static void
@@ -1532,6 +1531,13 @@ ary_deconstruct(mrb_state *mrb, mrb_value ary)
return ary;
}
enum {
comb_finished = 0,
comb_repeated_permutation = 1,
comb_repeated_combination = 2,
comb_permutation = 3,
comb_combination = 4
};
/*
* Internal method to initialize combination state.
@@ -1541,9 +1547,9 @@ static mrb_value
ary_combination_init(mrb_state *mrb, mrb_value self)
{
mrb_int k;
mrb_bool permutation;
mrb_sym mode_sym;
mrb_get_args(mrb, "ib", &k, &permutation);
mrb_get_args(mrb, "ni", &mode_sym, &k);
#if MRB_INT_MAX > SIZE_MAX
if (k > SIZE_MAX) {
mrb_raise(mrb, E_ARGUMENT_ERROR, "number too large");
@@ -1554,6 +1560,30 @@ ary_combination_init(mrb_state *mrb, mrb_value self)
return mrb_nil_value();
}
int mode;
switch (mode_sym) {
case MRB_SYM(repeated_permutation):
mode = comb_repeated_permutation;
break;
case MRB_SYM(repeated_combination):
mode = comb_repeated_combination;
break;
case MRB_SYM(permutation):
if (k > RARRAY_LEN(self)) {
return mrb_nil_value();
}
mode = comb_permutation;
break;
case MRB_SYM(combination):
if (k > RARRAY_LEN(self)) {
return mrb_nil_value();
}
mode = comb_combination;
break;
default:
mrb_raise(mrb, E_ARGUMENT_ERROR, "wrong mode");
}
struct RData *d;
struct mrb_combination_state *state;
Data_Make_Struct(mrb, mrb->object_class, struct mrb_combination_state,
@@ -1561,13 +1591,29 @@ ary_combination_init(mrb_state *mrb, mrb_value self)
state->k = k;
state->n = RARRAY_LEN(self);
state->permutation = permutation;
state->finished = FALSE;
state->mode = mode;
state->indices = (mrb_int*)mrb_calloc(mrb, k, sizeof(mrb_int));
if (mode == comb_permutation || mode == comb_combination) {
for (mrb_int i = 0; i < k; i++) {
state->indices[i] = i;
}
}
return mrb_obj_value(d);
}
static void
adjust_next_permutation_index(struct mrb_combination_state *state, mrb_int i)
{
for (mrb_int j = i - 1; j >= 0; j--) {
if (state->indices[i] == state->indices[j]) {
state->indices[i]++;
j = i;
}
}
}
/*
* Internal method to get next combination as index array.
* Returns array of indices or nil when iteration is complete.
@@ -1579,7 +1625,7 @@ ary_combination_next(mrb_state *mrb, mrb_value self)
mrb_get_args(mrb, "d", &state, &mrb_combination_state_type);
/* Check if iteration is complete */
if (state->finished) return mrb_nil_value();
if (state->mode == comb_finished) return mrb_nil_value();
/* Validate array hasn't been modified during iteration */
if (RARRAY_LEN(self) != state->n) {
@@ -1589,7 +1635,7 @@ ary_combination_next(mrb_state *mrb, mrb_value self)
/* Validate current indices are still in bounds */
for (mrb_int i = 0; i < state->k; i++) {
if (state->indices[i] >= state->n) {
state->finished = TRUE;
state->mode = comb_finished;
return mrb_nil_value();
}
}
@@ -1601,25 +1647,57 @@ ary_combination_next(mrb_state *mrb, mrb_value self)
mrb_ary_push(mrb, result, p[state->indices[i]]);
}
mrb_int pos = state->k - 1;
while (pos >= 0) {
state->indices[pos]++;
if (state->indices[pos] < state->n) break;
pos--;
}
if (pos < 0) {
state->finished = TRUE;
}
else {
/* Reset dependent indices */
mrb_int reset = state->permutation ? 0 : state->indices[pos];
for (pos++; pos < state->k; pos++) {
state->indices[pos] = reset;
switch (state->mode) {
case comb_repeated_permutation:
case comb_repeated_combination:
for (mrb_int i = state->k - 1; i >= 0; i--) {
state->indices[i]++;
if (state->indices[i] < state->n) {
/* Reset dependent indices */
mrb_int reset = (state->mode == comb_repeated_permutation) ? 0 : state->indices[i];
for (i++; i < state->k; i++) {
state->indices[i] = reset;
}
return result;
}
}
break;
case comb_permutation:
for (mrb_int i = state->k - 1; i >= 0; i--) {
state->indices[i]++;
// adjust so that it does not overlap with the leading index
adjust_next_permutation_index(state, i);
if (state->indices[i] < state->n) {
// adjust all trailing indexes to complete the function
for (i++; i < state->k; i++) {
state->indices[i] = 0;
adjust_next_permutation_index(state, i);
}
return result;
}
}
break;
case comb_combination:
for (mrb_int i = state->k - 1; i >= 0; i--) {
state->indices[i]++;
if (state->indices[i] <= state->n - state->k + i) {
// replace each overflowed indices with an index incremented by 1 from the previous one
for (i++; i < state->k; i++) {
state->indices[i] = state->indices[i - 1] + 1;
}
return result;
}
}
break;
default: // it probably wont happen, but just in case
result = mrb_nil_value();
break;
}
state->mode = comb_finished;
return result;
}