arrays: add binary_search, lower_bound, and upper_bound array operations (#12045)
parent
117091452b
commit
f282e64fe3
|
@ -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('')
|
||||
}
|
||||
|
|
|
@ -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
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue