From: Tanaka Akira Date: 2011-01-01T10:49:22+09:00 Subject: [ruby-dev:42915] min(n), max(n), min_by(n), max_by(n) 思ったんですが、 Enumerable#min Enumerable#max Enumerable#min_by Enumerable#max_by に、いくつとりだすか指定できる引数をつけるのはどうでしょうか。 p [3,1,4,1,5,9,2,6,5,3,5].min(3) #=> [1, 1, 2] p [3,1,4,1,5,9,2,6,5,3,5].max(3) #=> [5, 6, 9] first や last も引数をとるようになっていますから、それと似たかんじです。 max_by は [ruby-dev:42844] で使った Depq.nlargest のかわりに使えます。 % ./ruby -e ' module Enumerable def sample2(n) self.max_by(n) {|e| rand ** (1.0/(yield(e))) } end end enum = (-20..20).to_a*10000 a = enum.sample2(20000) {|e| Math.exp(-(e/5.0)**2) } h = Hash.new(0) a.each {|e| h[e] += 1 } -10.upto(10) {|x| puts "*" * (h[x]/30.0).to_i } ' * *** ****** *********** ****************** *************************** **************************************** *************************************************** ************************************************************* ******************************************************************** ************************************************************************ ************************************************************************ **************************************************************** **************************************************** ***************************************** **************************** ****************** ************ ****** *** * 中身は quick sort をベースにしたやりかたを使っています。 Index: enum.c =================================================================== --- enum.c (revision 30439) +++ enum.c (working copy) @@ -1104,6 +1104,176 @@ enum_none(VALUE obj) return result; } +struct nmin_data { + long n; + long bufmax; + long curlen; + VALUE buf; + int (*cmpfunc)(const void *, const void *, void *); + int rev; /* max if 1 */ + int by; /* min_by if 1 */ + const char *method; +}; + +static int +nmin_cmp(const void *ap, const void *bp, void *_data) +{ + struct nmin_data *data = (struct nmin_data *)_data; + VALUE a = *(const VALUE *)ap, b = *(const VALUE *)bp; + VALUE cmp = rb_funcall(a, id_cmp, 1, b); + if (RBASIC(data->buf)->klass) { + rb_raise(rb_eRuntimeError, "%s reentered", data->method); + } + return rb_cmpint(cmp, a, b); +} + +static int +nmin_block_cmp(const void *ap, const void *bp, void *_data) +{ + struct nmin_data *data = (struct nmin_data *)_data; + VALUE a = *(const VALUE *)ap, b = *(const VALUE *)bp; + VALUE cmp = rb_yield_values(2, a, b); + if (RBASIC(data->buf)->klass) { + rb_raise(rb_eRuntimeError, "%s reentered", data->method); + } + return rb_cmpint(cmp, a, b); +} + + +static void +nmin_filter(struct nmin_data *data) +{ + long n; + VALUE *beg; + int eltsize; + long numelts; + + long left, right; + + long i; + + if (data->curlen <= data->n) + return; + + n = data->n; + beg = RARRAY_PTR(data->buf); + eltsize = data->by ? 2 : 1; + numelts = data->curlen; + + left = 0; + right = numelts-1; + +#define GETPTR(i) (beg+(i)*eltsize) + +#define SWAP(i, j) do { \ + VALUE tmp[2]; \ + memcpy(tmp, GETPTR(i), sizeof(VALUE)*eltsize); \ + memcpy(GETPTR(i), GETPTR(j), sizeof(VALUE)*eltsize); \ + memcpy(GETPTR(j), tmp, sizeof(VALUE)*eltsize); \ +} while (0) + + while (1) { + long pivot_index = left + (right-left)/2; + long store_index; + + SWAP(pivot_index, right); + pivot_index = right; + + store_index = left; + for (i = left; i < right; i++) { + int c = data->cmpfunc(GETPTR(i), GETPTR(pivot_index), data); + if (data->rev) + c = -c; + if (c <= 0) { + SWAP(i, store_index); + store_index++; + } + } + SWAP(store_index, right); + + if (store_index == n || store_index == n-1) + break; + + if (n < store_index) { + right = store_index-1; + } + else { + left = store_index+1; + } + } +#undef GETPTR +#undef SWAP + + data->curlen = data->n; + ary_cutoff(data->buf, data->n * eltsize); +} + +static VALUE +nmin_i(VALUE i, VALUE *_data, int argc, VALUE *argv) +{ + struct nmin_data *data = (struct nmin_data *)_data; + + ENUM_WANT_SVALUE(); + + if (data->by) + rb_ary_push(data->buf, rb_yield(i)); + rb_ary_push(data->buf, i); + + data->curlen++; + + if (data->curlen == data->bufmax) { + nmin_filter(data); + } + + return Qnil; +} + +static VALUE +nmin_run(VALUE obj, VALUE num, int by, int rev) +{ + VALUE result; + struct nmin_data data; + + data.n = NUM2LONG(num); + if (data.n < 0) + rb_raise(rb_eArgError, "negative size (%ld)", data.n); + if (data.n == 0) + return rb_ary_new2(0); + if (LONG_MAX/4/(by ? 2 : 1) < data.n) + rb_raise(rb_eArgError, "too big size"); + data.bufmax = data.n * 4; + data.curlen = 0; + data.buf = rb_ary_tmp_new(data.bufmax * (by ? 2 : 1)); + data.cmpfunc = by ? nmin_cmp : + rb_block_given_p() ? nmin_block_cmp : + nmin_cmp; + data.rev = rev; + data.by = by; + data.method = rev ? (by ? "max_by" : "max") + : (by ? "min_by" : "min"); + rb_block_call(obj, id_each, 0, 0, nmin_i, (VALUE)&data); + nmin_filter(&data); + result = data.buf; + if (by) { + long i; + ruby_qsort(RARRAY_PTR(result), + RARRAY_LEN(result)/2, + sizeof(VALUE)*2, + data.cmpfunc, (void *)&data); + for (i=1; iklass = rb_cArray; + return result; + +} + static VALUE min_i(VALUE i, VALUE *memo, int argc, VALUE *argv) { @@ -1145,8 +1315,10 @@ min_ii(VALUE i, VALUE *memo, int argc, V /* * call-seq: - * enum.min -> obj - * enum.min {| a,b | block } -> obj + * enum.min -> obj + * enum.min {| a,b | block } -> obj + * enum.min(n) -> array + * enum.min(n) {| a,b | block } -> array * * Returns the object in enum with the minimum value. The * first form assumes all objects implement Comparable; @@ -1155,12 +1327,25 @@ min_ii(VALUE i, VALUE *memo, int argc, V * a = %w(albatross dog horse) * a.min #=> "albatross" * a.min {|a,b| a.length <=> b.length } #=> "dog" + * + * If the +n+ argument is given, minimum +n+ elements are returned + * as an array. + * + * a = %w[albatross dog horse] + * a.min(2) #=> ["albatross", "dog"] + * a.min(2) {|a, b| a.length <=> b.length } #=> ["dog", "horse"] */ static VALUE -enum_min(VALUE obj) +enum_min(int argc, VALUE *argv, VALUE obj) { VALUE result = Qundef; + VALUE num; + + rb_scan_args(argc, argv, "01", &num); + + if (!NIL_P(num)) + return nmin_run(obj, num, 0, 0); if (rb_block_given_p()) { rb_block_call(obj, id_each, 0, 0, min_ii, (VALUE)&result); @@ -1214,6 +1399,8 @@ max_ii(VALUE i, VALUE *memo, int argc, V * call-seq: * enum.max -> obj * enum.max {|a,b| block } -> obj + * enum.max(n) -> obj + * enum.max(n) {|a,b| block } -> obj * * Returns the object in _enum_ with the maximum value. The * first form assumes all objects implement Comparable; @@ -1222,12 +1409,25 @@ max_ii(VALUE i, VALUE *memo, int argc, V * a = %w(albatross dog horse) * a.max #=> "horse" * a.max {|a,b| a.length <=> b.length } #=> "albatross" + * + * If the +n+ argument is given, maximum +n+ elements are returned + * as an array. + * + * a = %w[albatross dog horse] + * a.max(2) #=> ["dog", "horse"] + * a.max(2) {|a, b| a.length <=> b.length } #=> ["horse", "albatross"] */ static VALUE -enum_max(VALUE obj) +enum_max(int argc, VALUE *argv, VALUE obj) { VALUE result = Qundef; + VALUE num; + + rb_scan_args(argc, argv, "01", &num); + + if (!NIL_P(num)) + return nmin_run(obj, num, 0, 1); if (rb_block_given_p()) { rb_block_call(obj, id_each, 0, 0, max_ii, (VALUE)&result); @@ -1410,8 +1610,10 @@ min_by_i(VALUE i, VALUE *memo, int argc, /* * call-seq: - * enum.min_by {|obj| block } -> obj - * enum.min_by -> an_enumerator + * enum.min_by {|obj| block } -> obj + * enum.min_by -> an_enumerator + * enum.min_by(n) {|obj| block } -> array + * enum.min_by(n) -> an_enumerator * * Returns the object in enum that gives the minimum * value from the given block. @@ -1420,14 +1622,26 @@ min_by_i(VALUE i, VALUE *memo, int argc, * * a = %w(albatross dog horse) * a.min_by {|x| x.length } #=> "dog" + * + * If the +n+ argument is given, minimum +n+ elements are returned + * as an array. + * + * a = %w[albatross dog horse] + * p a.min_by(2) {|x| x.length } #=> ["dog", "horse"] */ static VALUE -enum_min_by(VALUE obj) +enum_min_by(int argc, VALUE *argv, VALUE obj) { VALUE memo[2]; + VALUE num; - RETURN_ENUMERATOR(obj, 0, 0); + rb_scan_args(argc, argv, "01", &num); + + RETURN_ENUMERATOR(obj, argc, argv); + + if (!NIL_P(num)) + return nmin_run(obj, num, 1, 0); memo[0] = Qundef; memo[1] = Qnil; @@ -1456,8 +1670,10 @@ max_by_i(VALUE i, VALUE *memo, int argc, /* * call-seq: - * enum.max_by {|obj| block } -> obj - * enum.max_by -> an_enumerator + * enum.max_by {|obj| block } -> obj + * enum.max_by -> an_enumerator + * enum.max_by(n) {|obj| block } -> obj + * enum.max_by(n) -> an_enumerator * * Returns the object in enum that gives the maximum * value from the given block. @@ -1466,14 +1682,26 @@ max_by_i(VALUE i, VALUE *memo, int argc, * * a = %w(albatross dog horse) * a.max_by {|x| x.length } #=> "albatross" + * + * If the +n+ argument is given, minimum +n+ elements are returned + * as an array. + * + * a = %w[albatross dog horse] + * a.max_by(2) {|x| x.length } #=> ["horse", "albatross"] */ static VALUE -enum_max_by(VALUE obj) +enum_max_by(int argc, VALUE *argv, VALUE obj) { VALUE memo[2]; + VALUE num; - RETURN_ENUMERATOR(obj, 0, 0); + rb_scan_args(argc, argv, "01", &num); + + RETURN_ENUMERATOR(obj, argc, argv); + + if (!NIL_P(num)) + return nmin_run(obj, num, 1, 1); memo[0] = Qundef; memo[1] = Qnil; @@ -2692,11 +2920,11 @@ Init_Enumerable(void) rb_define_method(rb_mEnumerable, "any?", enum_any, 0); rb_define_method(rb_mEnumerable, "one?", enum_one, 0); rb_define_method(rb_mEnumerable, "none?", enum_none, 0); - rb_define_method(rb_mEnumerable, "min", enum_min, 0); - rb_define_method(rb_mEnumerable, "max", enum_max, 0); + rb_define_method(rb_mEnumerable, "min", enum_min, -1); + rb_define_method(rb_mEnumerable, "max", enum_max, -1); rb_define_method(rb_mEnumerable, "minmax", enum_minmax, 0); - rb_define_method(rb_mEnumerable, "min_by", enum_min_by, 0); - rb_define_method(rb_mEnumerable, "max_by", enum_max_by, 0); + rb_define_method(rb_mEnumerable, "min_by", enum_min_by, -1); + rb_define_method(rb_mEnumerable, "max_by", enum_max_by, -1); rb_define_method(rb_mEnumerable, "minmax_by", enum_minmax_by, 0); rb_define_method(rb_mEnumerable, "member?", enum_member, 1); rb_define_method(rb_mEnumerable, "include?", enum_member, 1); Index: range.c =================================================================== --- range.c (revision 30439) +++ range.c (working copy) @@ -612,8 +612,10 @@ range_last(int argc, VALUE *argv, VALUE /* * call-seq: - * rng.min -> obj - * rng.min {| a,b | block } -> obj + * rng.min -> obj + * rng.min {| a,b | block } -> obj + * rng.min(n) -> array + * rng.min(n) {| a,b | block } -> array * * Returns the minimum value in rng. The second uses * the block to compare values. Returns nil if the first @@ -623,10 +625,13 @@ range_last(int argc, VALUE *argv, VALUE static VALUE -range_min(VALUE range) +range_min(int argc, VALUE *argv, VALUE range) { if (rb_block_given_p()) { - return rb_call_super(0, 0); + return rb_call_super(argc, argv); + } + else if (argc != 0) { + return range_first(argc, argv, range); } else { VALUE b = RANGE_BEG(range); @@ -641,8 +646,10 @@ range_min(VALUE range) /* * call-seq: - * rng.max -> obj - * rng.max {| a,b | block } -> obj + * rng.max -> obj + * rng.max {| a,b | block } -> obj + * rng.max(n) -> obj + * rng.max(n) {| a,b | block } -> obj * * Returns the maximum value in rng. The second uses * the block to compare values. Returns nil if the first @@ -651,13 +658,13 @@ range_min(VALUE range) */ static VALUE -range_max(VALUE range) +range_max(int argc, VALUE *argv, VALUE range) { VALUE e = RANGE_END(range); int nm = FIXNUM_P(e) || rb_obj_is_kind_of(e, rb_cNumeric); - if (rb_block_given_p() || (EXCL(range) && !nm)) { - return rb_call_super(0, 0); + if (rb_block_given_p() || (EXCL(range) && !nm) || argc) { + return rb_call_super(argc, argv); } else { VALUE b = RANGE_BEG(range); @@ -1041,8 +1048,8 @@ Init_Range(void) rb_define_method(rb_cRange, "end", range_end, 0); rb_define_method(rb_cRange, "first", range_first, -1); rb_define_method(rb_cRange, "last", range_last, -1); - rb_define_method(rb_cRange, "min", range_min, 0); - rb_define_method(rb_cRange, "max", range_max, 0); + rb_define_method(rb_cRange, "min", range_min, -1); + rb_define_method(rb_cRange, "max", range_max, -1); rb_define_method(rb_cRange, "to_s", range_to_s, 0); rb_define_method(rb_cRange, "inspect", range_inspect, 0); Index: test/ruby/test_range.rb =================================================================== --- test/ruby/test_range.rb (revision 30439) +++ test/ruby/test_range.rb (working copy) @@ -58,6 +58,9 @@ class TestRange < Test::Unit::TestCase assert_equal(0, (0..0).min) assert_equal(nil, (0...0).min) + + assert_equal([0,1,2], (0..10).min(3)) + assert_equal([0,1], (0..1).min(3)) end def test_max @@ -73,6 +76,9 @@ class TestRange < Test::Unit::TestCase assert_equal(0, (0..0).max) assert_equal(nil, (0...0).max) + + assert_equal([8,9,10], (0..10).max(3)) + assert_equal([7,8,9], (0...10).max(3)) end def test_initialize_twice Index: test/ruby/test_enum.rb =================================================================== --- test/ruby/test_enum.rb (revision 30439) +++ test/ruby/test_enum.rb (working copy) @@ -147,6 +147,9 @@ class TestEnumerable < Test::Unit::TestC assert_equal("albatross", ary.min) assert_equal("dog", ary.min {|a,b| a.length <=> b.length }) assert_equal(1, [3,2,1].min) + assert_equal(%w[albatross dog], ary.min(2)) + assert_equal(%w[dog horse], + ary.min(2) {|a,b| a.length <=> b.length }) end def test_max @@ -156,6 +159,9 @@ class TestEnumerable < Test::Unit::TestC assert_equal("horse", ary.max) assert_equal("albatross", ary.max {|a,b| a.length <=> b.length }) assert_equal(1, [3,2,1].max{|a,b| b <=> a }) + assert_equal(%w[dog horse], ary.max(2)) + assert_equal(%w[horse albatross], + ary.max(2) {|a,b| a.length <=> b.length }) end def test_minmax @@ -174,6 +180,7 @@ class TestEnumerable < Test::Unit::TestC a = %w(albatross dog horse) assert_equal("dog", a.min_by {|x| x.length }) assert_equal(3, [2,3,1].min_by {|x| -x }) + assert_equal(%w[dog horse], a.min_by(2) {|x| x.length }) end def test_max_by @@ -181,6 +188,7 @@ class TestEnumerable < Test::Unit::TestC a = %w(albatross dog horse) assert_equal("albatross", a.max_by {|x| x.length }) assert_equal(1, [2,3,1].max_by {|x| -x }) + assert_equal(%w[horse albatross], a.max_by(2) {|x| x.length }) end def test_minmax_by -- [田中 哲][たなか あきら][Tanaka Akira]