Ruby » Ruby master

/*

* call-seq:

* ary.sample! -> obj

* ary.sample!(random: rng) -> obj

* ary.sample!(n) -> new_ary

* ary.sample!(n, random: rng) -> new_ary

*

* Choose and delete 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

* <code>nil</code> 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)

{

VALUE nv, result, *ptr_result;

VALUE opts, randgen = rb_cRandom;

VALUE a, b, c;

long n, len, i, j, k;

double rnds[3];

if (OPTHASH_GIVEN_P(opts)) {

randgen = rb_hash_lookup2(opts, sym_random, randgen);

}

len = RARRAY_LEN(ary);

if (argc == 0) {

if (len == 0) return Qnil;

if (len == 1) {

i = 0;

}

else {

double x = rb_random_real(randgen);

if (len == 0) return Qnil;

i = (long)(x * len);

}

return rb_ary_delete_at(ary, i);

}

rb_scan_args(argc, argv, "1", &nv);

n = NUM2LONG(nv);

if (n < 0) rb_raise(rb_eArgError, "negative sample number");

if (n > len) n = len;

if (n <= numberof(rnds)) {

for (i = 0; i < n; ++i) {

rnds[i] = rb_random_real(randgen);

}

}

switch (n) {

case 0:

return rb_ary_new2(0);

case 1:

i = (long)(rnds[0] * len);

return rb_ary_new3(1, rb_ary_delete_at(ary, i));

case 2:

i = (long)(rnds[0] * len);

j = (long)(rnds[1] * (len-1));

a = rb_ary_delete_at(ary, i);

b = rb_ary_delete_at(ary, j);

return rb_ary_new3(2, a, b);

case 3:

i = (long)(rnds[0] * len);

j = (long)(rnds[1] * (len-1));

k = (long)(rnds[2] * (len-2));

a = rb_ary_delete_at(ary, i);

b = rb_ary_delete_at(ary, j);

c = rb_ary_delete_at(ary, k);

return rb_ary_new3(3, a, b, c);

}

result = rb_ary_new2(n);

ptr_result = RARRAY_PTR(result);

for (i=0; i<n; i++) {

j = (long)(rb_random_real(randgen) * --len);

ptr_result[i] = rb_ary_delete_at(ary, j);

}

ARY_SET_LEN(result, n);

return result;

}

rb_define_method(rb_cArray, "sample!", rb_ary_sample_bang, -1);