From f282e64fe345a567051052e4b40e50c744363020 Mon Sep 17 00:00:00 2001
From: 05st <60903484+05st@users.noreply.github.com>
Date: Sun, 3 Oct 2021 00:14:39 -0500
Subject: [PATCH] arrays: add binary_search, lower_bound, and upper_bound array
 operations (#12045)

---
 vlib/arrays/arrays.v      | 63 +++++++++++++++++++++++++++++++++++++++
 vlib/arrays/arrays_test.v | 24 +++++++++++++++
 2 files changed, 87 insertions(+)

diff --git a/vlib/arrays/arrays.v b/vlib/arrays/arrays.v
index 1dd3310755..60bd8e8dc0 100644
--- a/vlib/arrays/arrays.v
+++ b/vlib/arrays/arrays.v
@@ -290,3 +290,66 @@ pub fn concat<T>(a []T, b ...T) []T {
 
 	return m
 }
+
+// returns the smallest element >= val, requires `arr` to be sorted
+pub fn lower_bound<T>(arr []T, val T) ?T {
+	if arr.len == 0 {
+		return error('.lower_bound called on an empty array')
+	}
+	mut left, mut right := 0, arr.len - 1
+	for ; left <= right; {
+		idx := (left + right) / 2
+		elem := arr[idx]
+		if elem < val {
+			left = idx + 1
+		} else {
+			right = idx - 1
+		}
+	}
+	if left >= arr.len {
+		return error('')
+	} else {
+		return arr[left]
+	}
+}
+
+// returns the largest element <= val, requires `arr` to be sorted
+pub fn upper_bound<T>(arr []T, val T) ?T {
+	if arr.len == 0 {
+		return error('.upper_bound called on an empty array')
+	}
+	mut left, mut right := 0, arr.len - 1
+	for ; left <= right; {
+		idx := (left + right) / 2
+		elem := arr[idx]
+		if elem > val {
+			right = idx - 1
+		} else {
+			left = idx + 1
+		}
+	}
+	if right < 0 {
+		return error('')
+	} else {
+		return arr[right]
+	}
+}
+
+// binary search, requires `arr` to be sorted, returns index
+pub fn binary_search<T>(arr []T, target T) ?int {
+	mut left := 0
+	mut right := arr.len - 1
+	for ; left <= right; {
+		idx := (left + right) / 2
+		elem := arr[idx]
+		if elem == target {
+			return idx
+		}
+		if elem < target {
+			left = idx + 1
+		} else {
+			right = idx - 1
+		}
+	}
+	return error('')
+}
diff --git a/vlib/arrays/arrays_test.v b/vlib/arrays/arrays_test.v
index 29cc248b2a..b5563bc87a 100644
--- a/vlib/arrays/arrays_test.v
+++ b/vlib/arrays/arrays_test.v
@@ -194,3 +194,27 @@ fn test_concat_string() {
 
 	assert concat(a, ...b) == ['1', '2', '3', '3', '2', '1']
 }
+
+fn test_binary_search() ? {
+	a := [1, 3, 3, 4, 5, 6, 7, 8, 10]
+	assert binary_search(a, 3) ? == 1
+	assert (binary_search(a, 0) or { -1 }) == -1
+}
+
+fn test_lower_bound() ? {
+	a := [1, 3, 3, 4, 5, 6, 7, 8, 10]
+	b := []int{}
+	c := [1, 2, 3]
+	assert lower_bound(a, 2) ? == 3
+	assert (lower_bound(b, 4) or { -1 }) == -1
+	assert lower_bound(c, 3) ? == 3
+}
+
+fn test_upper_bound() ? {
+	a := [1, 3, 3, 4, 5, 6, 7, 8, 10]
+	b := []int{}
+	c := [1, 2, 3]
+	assert upper_bound(a, 9) ? == 8
+	assert (upper_bound(b, 4) or { -1 }) == -1
+	assert upper_bound(c, 2) ? == 2
+}