diff --git a/array.c b/array.c index a19e7bb..0c0e3fb 100755 --- a/array.c +++ b/array.c @@ -4766,6 +4766,51 @@ rb_ary_shuffle(int argc, VALUE *argv, VALUE ary) } +struct sample_with_memo_arg { + VALUE randgen, ary; + long len, n; + st_table *memo; + VALUE result; +}; + +static VALUE +sample_with_memo(VALUE a) +{ + struct sample_with_memo_arg *arg = (void *)a; + VALUE randgen = arg->randgen; + long i, j, i2, j2, len; + st_data_t value; + RARRAY_PTR_USE(arg->result, ptr_result, { + for (i=0; in; i++) { + ptr_result[i] = RAND_UPTO(arg->len-i) + i; + } + len = RARRAY_LEN(arg->ary); + if (len < arg->len) { + for (i=0; in; i++) { + if ((long)ptr_result[i] >= len) arg->n = 0; + } + if (arg->n > len) arg->n = len; + } + for (i=0; in; i++) { + j2 = j = ptr_result[i]; + i2 = i; + if (st_lookup(arg->memo, (st_data_t)i, &value)) i2 = (long)value; + if (st_lookup(arg->memo, (st_data_t)j, &value)) j2 = (long)value; + st_insert(arg->memo, (st_data_t)j, (st_data_t)i2); + ptr_result[i] = rb_ary_elt(arg->ary, j2); + } + }); + return Qnil; +} + +static VALUE +ensure_sample_with_memo(VALUE arg) +{ + st_free_table((st_table*)arg); + return Qnil; +} + + /* * call-seq: * ary.sample -> obj @@ -4875,6 +4920,13 @@ rb_ary_sample(int argc, VALUE *argv, VALUE ary) } }); } + else if (n <= (len<2560 ? len/128 : len<5120 ? len/64 : len<10240 ? len/32 : len/16)) { + struct sample_with_memo_arg arg = { randgen, ary, len, n }; + arg.memo = st_init_numtable_with_size(n); + result = arg.result = rb_ary_new_capa(n); + rb_ensure(sample_with_memo, (VALUE)&arg, ensure_sample_with_memo, (VALUE)arg.memo); + n = arg.n; + } else { result = rb_ary_dup(ary); RBASIC_CLEAR_CLASS(result);