arrays: add binary_search, lower_bound, and upper_bound array operations (#12045)

pull/12050/head
05st 2021-10-03 00:14:39 -05:00 committed by GitHub
parent 117091452b
commit f282e64fe3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 87 additions and 0 deletions

View File

@ -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('')
}

View File

@ -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
}