Add Array#bsearch

2 files changed, 98 insertions, 0 deletions

diff --git a/mrbgems/mruby-array-ext/mrblib/array.rb b/mrbgems/mruby-array-ext/mrblib/array.rb
index 49d0db0d5..411c4e89f 100644
--- a/mrbgems/mruby-array-ext/mrblib/array.rb
+++ b/mrbgems/mruby-array-ext/mrblib/array.rb
@@ -509,4 +509,86 @@ class Array
     self[idx, 0] = args
     self
   end
+
+  ##
+  #  call-seq:
+  #     ary.bsearch {|x| block }  -> elem
+  #
+  #  By using binary search, finds a value from this array which meets
+  #  the given condition in O(log n) where n is the size of the array.
+  #
+  #  You can use this method in two use cases: a find-minimum mode and
+  #  a find-any mode.  In either case, the elements of the array must be
+  #  monotone (or sorted) with respect to the block.
+  #
+  #  In find-minimum mode (this is a good choice for typical use case),
+  #  the block must return true or false, and there must be an index i
+  #  (0 <= i <= ary.size) so that:
+  #
+  #  - the block returns false for any element whose index is less than
+  #    i, and
+  #  - the block returns true for any element whose index is greater
+  #    than or equal to i.
+  #
+  #  This method returns the i-th element.  If i is equal to ary.size,
+  #  it returns nil.
+  #
+  #     ary = [0, 4, 7, 10, 12]
+  #     ary.bsearch {|x| x >=   4 } #=> 4
+  #     ary.bsearch {|x| x >=   6 } #=> 7
+  #     ary.bsearch {|x| x >=  -1 } #=> 0
+  #     ary.bsearch {|x| x >= 100 } #=> nil
+  #
+  #  In find-any mode (this behaves like libc's bsearch(3)), the block
+  #  must return a number, and there must be two indices i and j
+  #  (0 <= i <= j <= ary.size) so that:
+  #
+  #  - the block returns a positive number for ary[k] if 0 <= k < i,
+  #  - the block returns zero for ary[k] if i <= k < j, and
+  #  - the block returns a negative number for ary[k] if
+  #    j <= k < ary.size.
+  #
+  #  Under this condition, this method returns any element whose index
+  #  is within i...j.  If i is equal to j (i.e., there is no element
+  #  that satisfies the block), this method returns nil.
+  #
+  #     ary = [0, 4, 7, 10, 12]
+  #     # try to find v such that 4 <= v < 8
+  #     ary.bsearch {|x| 1 - (x / 4).truncate } #=> 4 or 7
+  #     # try to find v such that 8 <= v < 10
+  #     ary.bsearch {|x| 4 - (x / 2).truncate } #=> nil
+  #
+  #  You must not mix the two modes at a time; the block must always
+  #  return either true/false, or always return a number.  It is
+  #  undefined which value is actually picked up at each iteration.
+
+  def bsearch(&block)
+    return to_enum :bsearch unless block_given?
+
+    low = 0
+    high = self.size
+    satisfied = false
+    while low < high
+      mid = low + ((high - low) / 2).truncate
+      val = self[mid]
+      v = block.call(val)
+      if v.is_a?(Integer)
+        return val if v == 0
+        smaller = v < 0
+      elsif v == true
+        satisfied = true
+        smaller = true
+      elsif v == false || v == nil
+        smaller = false
+      end
+      if smaller
+        high = mid
+      else
+        low = mid + 1
+      end
+    end
+    return nil if low == self.size
+    return nil unless satisfied
+    self[low]
+  end
 end
diff --git a/mrbgems/mruby-array-ext/test/array.rb b/mrbgems/mruby-array-ext/test/array.rb
index d15ea2a64..ff82b7bfa 100644
--- a/mrbgems/mruby-array-ext/test/array.rb
+++ b/mrbgems/mruby-array-ext/test/array.rb
@@ -206,3 +206,19 @@ assert("Array#insert") do
   b = ["a", "b", "c", "d"]
   assert_equal ["a", "b", "c", "d", nil, nil, 99], b.insert(6, 99)
 end
+
+assert("Array#bsearch") do
+  # Find minimum mode
+  a = [0, 4, 7, 10, 12]
+  assert_include [4, 7], a.bsearch {|x| x >= 4 }
+  assert_equal 7, a.bsearch {|x| x >= 6 }
+  assert_equal 0, a.bsearch {|x| x >= -1 }
+  assert_nil a.bsearch {|x| x >= 100 }
+
+  # Find any mode
+  a = [0, 4, 7, 10, 12]
+  assert_include [4, 7], a.bsearch {|x| 1 - (x / 4).truncate }
+  assert_nil a.bsearch {|x| 4 - (x / 2).truncate }
+  assert_equal(nil, a.bsearch {|x| 1 })
+  assert_equal(nil, a.bsearch {|x| -1 })
+end