diff --git a/array.c b/array.c index cfe44a9..392f366 100644 --- a/array.c +++ b/array.c @@ -3934,29 +3934,12 @@ rb_ary_shuffle(int argc, VALUE *argv, VALUE ary) return ary; } - -/* - * call-seq: - * ary.sample -> obj - * ary.sample(random: rng) -> obj - * ary.sample(n) -> new_ary - * ary.sample(n, random: rng) -> new_ary - * - * Choose a random element or +n+ random elements from the array. The elements - * are chosen by using random and unique indices into the array in order to - * ensure that an element doesn't repeat itself unless the array already - * contained duplicate elements. If the array is empty the first form returns - * nil and the second form returns an empty array. - * - * If +rng+ is given, it will be used as the random number generator. - */ - - static VALUE -rb_ary_sample(int argc, VALUE *argv, VALUE ary) +ary_sample(int argc, VALUE *argv, VALUE ary, int removing) { VALUE nv, result, *ptr; VALUE opts, randgen = rb_cRandom; + VALUE *ptr_result; long n, len, i, j, k, idx[10]; double rnds[numberof(idx)]; @@ -3975,7 +3958,7 @@ rb_ary_sample(int argc, VALUE *argv, VALUE ary) if ((len = RARRAY_LEN(ary)) == 0) return Qnil; i = (long)(x * len); } - return RARRAY_PTR(ary)[i]; + return (removing ? rb_ary_delete_at(ary, i) : RARRAY_PTR(ary)[i]); } rb_scan_args(argc, argv, "1", &nv); n = NUM2LONG(nv); @@ -3994,12 +3977,18 @@ rb_ary_sample(int argc, VALUE *argv, VALUE ary) return rb_ary_new2(0); case 1: i = (long)(rnds[0] * len); - return rb_ary_new4(1, &ptr[i]); + return (removing ? rb_ary_new3(1, rb_ary_delete_at(ary, i)) : rb_ary_new4(1, &ptr[i])); case 2: i = (long)(rnds[0] * len); j = (long)(rnds[1] * (len-1)); if (j >= i) j++; - return rb_ary_new3(2, ptr[i], ptr[j]); + result = rb_ary_new3(2, ptr[i], ptr[j]); + if (removing) { + rb_ary_delete_at(ary, i); + if (i < j) j--; + rb_ary_delete_at(ary, j); + } + return result; case 3: i = (long)(rnds[0] * len); j = (long)(rnds[1] * (len-1)); @@ -4009,10 +3998,26 @@ rb_ary_sample(int argc, VALUE *argv, VALUE ary) if (j >= i) l = i, g = ++j; if (k >= l && (++k >= g)) ++k; } - return rb_ary_new3(3, ptr[i], ptr[j], ptr[k]); + result = rb_ary_new3(3, ptr[i], ptr[j], ptr[k]); + if (removing) { + long l, g; + if (i < k) l = i, g = k, i = g, k = l; + if (i < j) l = i, g = j, i = g, j = l; + rb_ary_delete_at(ary, i); + rb_ary_delete_at(ary, j); + rb_ary_delete_at(ary, k); + } + return result; } - if (n <= numberof(idx)) { - VALUE *ptr_result; + if (removing) { + result = rb_ary_new2(n); + ptr_result = RARRAY_PTR(result); + for (i=0; iklass = 0; ptr_result = RARRAY_PTR(result); @@ -4049,6 +4053,49 @@ rb_ary_sample(int argc, VALUE *argv, VALUE ary) return result; } +/* + * call-seq: + * ary.sample -> obj + * ary.sample(random: rng) -> obj + * ary.sample(n) -> new_ary + * ary.sample(n, random: rng) -> new_ary + * + * Choose a random element or +n+ random elements from the array. The elements + * are chosen by using random and unique indices into the array in order to + * ensure that an element doesn't repeat itself unless the array already + * contained duplicate elements. If the array is empty the first form returns + * nil and the second form returns an empty array. + * + * If +rng+ is given, it will be used as the random number generator. + */ + +static VALUE +rb_ary_sample(int argc, VALUE *argv, VALUE ary) +{ + return ary_sample(argc, argv, ary, 0); +} + +/* + * call-seq: + * ary.sample! -> obj + * ary.sample!(random: rng) -> obj + * ary.sample!(n) -> new_ary + * ary.sample!(n, random: rng) -> new_ary + * + * Choose and remove a random element or +n+ random elements from the array. + * The elements are chosen by using random and unique indices into the array in + * order to ensure that an element doesn't repeat itself unless the array + * already contained duplicate elements. If the array is empty the first form + * returns nil and the second form returns an empty array. + * + * If +rng+ is given, it will be used as the random number generator. + */ + +static VALUE +rb_ary_sample_bang(int argc, VALUE *argv, VALUE ary) +{ + return ary_sample(argc, argv, ary, 1); +} /* * call-seq: @@ -4998,6 +5045,7 @@ Init_Array(void) rb_define_method(rb_cArray, "shuffle!", rb_ary_shuffle_bang, -1); rb_define_method(rb_cArray, "shuffle", rb_ary_shuffle, -1); rb_define_method(rb_cArray, "sample", rb_ary_sample, -1); + rb_define_method(rb_cArray, "sample!", rb_ary_sample_bang, -1); rb_define_method(rb_cArray, "cycle", rb_ary_cycle, -1); rb_define_method(rb_cArray, "permutation", rb_ary_permutation, -1); rb_define_method(rb_cArray, "combination", rb_ary_combination, 1);