[ruby-dev:34067] Array#take,take_while,drop,drop_whlie

遠藤と申します。

ary.drop(1) は each を経由するので常に O(n) かかり、ちょっと遅いです。

$ time ./ruby -e 'a=[0]*10000; a = a.drop(1) until a.empty?'

real    0m4.440s
user    0m4.310s
sys     0m0.090s

そこで、Array 特化版の Array#take 、take_while 、drop、drop_while を
作ってみました。バッファの共有がきくので速くなります。

$ time ./ruby -e 'a=[0]*10000; a = a.drop(1) until a.empty?'

real    0m0.210s
user    0m0.020s
sys     0m0.030s

どうでしょうか。


Index: array.c
===================================================================
--- array.c	(revision 15784)
+++ array.c	(working copy)
@@ -3228,8 +3228,100 @@
     return result;
 }

+/*
+ *  call-seq:
+ *     ary.take(n)               => array
+ *
+ *  Returns first n elements from <i>ary</i>.
+ *
+ *     a = [1, 2, 3, 4, 5, 0]
+ *     a.take(3)             # => [1, 2, 3]
+ *
+ */

+static VALUE
+rb_ary_take(VALUE obj, VALUE n)
+{
+    return rb_ary_subseq(obj, 0, FIX2LONG(n));
+}

+/*
+ *  call-seq:
+ *     ary.take_while {|arr| block }   => array
+ *
+ *  Passes elements to the block until the block returns nil or false,
+ *  then stops iterating and returns an array of all prior elements.
+ *
+ *     a = [1, 2, 3, 4, 5, 0]
+ *     a.take_while {|i| i < 3 }   # => [1, 2]
+ *
+ */
+
+static VALUE
+rb_ary_take_while(VALUE ary)
+{
+    VALUE result;
+    long i;
+
+    RETURN_ENUMERATOR(ary, 0, 0);
+    result = rb_ary_new2(RARRAY_LEN(ary));
+    for (i = 0; i < RARRAY_LEN(ary); i++) {
+	if (!RTEST(rb_yield(RARRAY_PTR(ary)[i]))) break;
+	rb_ary_push(result, rb_ary_elt(ary, i));
+    }
+    return result;
+}
+
+/*
+ *  call-seq:
+ *     ary.drop(n)               => array
+ *
+ *  Drops first n elements from <i>ary</i>, and returns rest elements
+ *  in an array.
+ *
+ *     a = [1, 2, 3, 4, 5, 0]
+ *     a.drop(3)             # => [4, 5, 0]
+ *
+ */
+
+static VALUE
+rb_ary_drop(VALUE ary, VALUE n)
+{
+    VALUE result;
+
+    result = rb_ary_subseq(ary, FIX2LONG(n), RARRAY_LEN(ary));
+    if (result == Qnil) result = rb_ary_new();
+    return result;
+}
+
+/*
+ *  call-seq:
+ *     ary.drop_while {|arr| block }   => array
+ *
+ *  Drops elements up to, but not including, the first element for
+ *  which the block returns nil or false and returns an array
+ *  containing the remaining elements.
+ *
+ *     a = [1, 2, 3, 4, 5, 0]
+ *     a.drop_while {|i| i < 3 }   # => [3, 4, 5, 0]
+ *
+ */
+
+static VALUE
+rb_ary_drop_while(VALUE ary)
+{
+    VALUE result;
+    long i;
+
+    RETURN_ENUMERATOR(ary, 0, 0);
+    for (i = 0; i < RARRAY_LEN(ary); i++) {
+	if (!RTEST(rb_yield(RARRAY_PTR(ary)[i]))) break;
+    }
+    return rb_ary_drop(ary, LONG2FIX(i));
+}
+
+
+
 /* Arrays are ordered, integer-indexed collections of any object.
  * Array indexing starts at 0, as in C or Java.  A negative index is
  * assumed to be relative to the end of the array---that is, an index of -1
@@ -3332,5 +3424,10 @@
     rb_define_method(rb_cArray, "combination", rb_ary_combination, 1);
     rb_define_method(rb_cArray, "product", rb_ary_product, -1);

+    rb_define_method(rb_cArray, "take", rb_ary_take, 1);
+    rb_define_method(rb_cArray, "take_while", rb_ary_take_while, 0);
+    rb_define_method(rb_cArray, "drop", rb_ary_drop, 1);
+    rb_define_method(rb_cArray, "drop_while", rb_ary_drop_while, 0);
+
     id_cmp = rb_intern("<=>");
 }
Index: test/ruby/test_array.rb
===================================================================
--- test/ruby/test_array.rb	(revision 15784)
+++ test/ruby/test_array.rb	(working copy)
@@ -1229,4 +1229,20 @@
     assert_equal(@cls[].permutation(0).to_a, @cls[[]])

   end
+
+  def test_take
+    assert_equal([1,2,3], [1,2,3,4,5,0].take(3))
+  end
+
+  def test_take_while
+    assert_equal([1,2], [1,2,3,4,5,0].take_while {|i| i < 3 })
+  end
+
+  def test_drop
+    assert_equal([4,5,0], [1,2,3,4,5,0].drop(3))
+  end
+
+  def test_drop_while
+    assert_equal([3,4,5,0], [1,2,3,4,5,0].drop_while {|i| i < 3 })
+  end
 end

-- 
Yusuke ENDOH <[e-mail]>