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
|
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']
|
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