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);