diff --git a/array.c b/array.c index cfe44a9..41fd555 100644 --- a/array.c +++ b/array.c @@ -4049,6 +4049,87 @@ 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 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 + * 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) +{ + 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