diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb index 5f74489f8..35abf95b4 100644 --- a/mrbgems/mruby-array-ext/mrblib/array.rb +++ b/mrbgems/mruby-array-ext/mrblib/array.rb @@ -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 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 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"); } @@ -1580,9 +1594,26 @@ ary_combination_init(mrb_state *mrb, mrb_value self) 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. @@ -1631,6 +1662,36 @@ ary_combination_next(mrb_state *mrb, mrb_value self) } } 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 won’t happen, but just in case result = mrb_nil_value(); break;