Add lowerBound/upperBound/equalRange

Authored by https://github.com/CraigglesO

Original Discussion: #9890

I already had to create these functions and test cases for my own
project so I decided to contribute to the main code base in hopes it
would simplify my own.

I realize this is still under discussion but this was a trivial amount
of work so I thought I could help nudge the discussion towards a
decision.

Why add these to the standard library

To better illustrate and solidify their value, the standard library's
"sort" module already contains several binary search queries on arrays
such as binarySearch and internal functions binaryFirst & binaryLast. A
final example of its use: the Zig code itself created and used a
bounding search in the linker.

There still lacks the ability to allow the programmer themselves to
search the array for a rough position and find an index to read &/
update.

Adding these functions would also help to complement dynamic structures
like ArrayList with it's insert function.

Example Case

I'm building a library in Zig for GIS geometry. To store points, lines,
and polygons each 3D point is first translated into what's called an
S2CellId. This is a fancy way of saying I reduce the Earth's spherical
data into a 1D Hilbert Curve with cm precision. This gives me 2
convenient truths:

    Hilbert Curves have locality of reference.
    All points can be stored inside a 1D array

Since lowerBound and upperBound to find data inside a radius for
instance. If I'm interested in a specific cell at a specific "level" and
want to iterate all duplicates, equalRange is a best fit.
This commit is contained in:
Craig O'Connor 2024-01-27 22:42:56 +01:00 committed by John Schmidt
parent ee36131e6a
commit 664c18544c

View File

@ -502,6 +502,290 @@ test "binarySearch" {
);
}
/// Returns the index pointing to the first element in the range [first,last)
/// which does not compare less than val.
/// The function optimizes the number of comparisons performed by using a binary search O(log n).
/// An example lessThan function:
/// fn lower_u32(_: void, key: u32, lhs: u32) bool { return lhs < key; }
pub fn lowerBound(
comptime T: type,
key: anytype,
items: []const T,
context: anytype,
comptime lessThan: fn (context: @TypeOf(context), key: @TypeOf(key), lhs: T) bool,
) usize {
var left: usize = 0;
var right: usize = items.len;
var mid: usize = undefined;
while (left < right) {
mid = left + (right - left) / 2;
if (lessThan(context, key, items[mid])) {
left = mid + 1;
} else {
right = mid;
}
}
return left;
}
test "lowerBound" {
const S = struct {
fn lower_u32(context: void, key: u32, lhs: u32) bool {
_ = context;
return lhs < key;
}
fn lower_i32(context: void, key: i32, lhs: i32) bool {
_ = context;
return lhs < key;
}
fn lower_f32(context: void, key: f32, lhs: f32) bool {
_ = context;
return lhs < key;
}
};
// u32
// test no data
try testing.expectEqual(
@as(?usize, 0),
lowerBound(u32, @as(u32, 0), &[_]u32{}, {}, S.lower_u32),
);
// test below the first element
try testing.expectEqual(
@as(?usize, 0),
lowerBound(u32, @as(u32, 0), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
);
// test equal to the first element
try testing.expectEqual(
@as(?usize, 0),
lowerBound(u32, @as(u32, 2), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
);
// test between two numbers
try testing.expectEqual(
@as(?usize, 2),
lowerBound(u32, @as(u32, 5), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
);
// test equal to a number not at the ends
try testing.expectEqual(
@as(?usize, 2),
lowerBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
);
// test equal to the last element
try testing.expectEqual(
@as(?usize, 5),
lowerBound(u32, @as(u32, 64), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
);
// test above the last element
try testing.expectEqual(
@as(?usize, 6),
lowerBound(u32, @as(u32, 100), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
);
// i32
try testing.expectEqual(
@as(?usize, 2),
lowerBound(i32, @as(i32, 5), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
);
// f32
try testing.expectEqual(
@as(?usize, 1),
lowerBound(f32, @as(f32, -33.4), &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, {}, S.lower_f32),
);
}
/// Returns the index pointing to the first element in the range [first,last)
/// which compares greater than val.
/// The function optimizes the number of comparisons performed by using a binary search O(log n).
/// An example greaterThan function:
/// fn upper_u32(_: void, key: u32, rhs: u32) bool { return key >= rhs; }
pub fn upperBound(
comptime T: type,
key: anytype,
items: []const T,
context: anytype,
comptime greaterThan: fn (context: @TypeOf(context), key: @TypeOf(key), rhs: T) bool,
) usize {
var left: usize = 0;
var right: usize = items.len;
var mid: usize = undefined;
while (left < right) {
mid = (right + left) / 2;
if (greaterThan(context, key, items[mid])) {
left = mid + 1;
} else {
right = mid;
}
}
return left;
}
test "upperBound" {
const S = struct {
fn upper_u32(context: void, key: u32, rhs: u32) bool {
_ = context;
return key >= rhs;
}
fn upper_i32(context: void, key: i32, rhs: i32) bool {
_ = context;
return key >= rhs;
}
fn upper_f32(context: void, key: f32, rhs: f32) bool {
_ = context;
return key >= rhs;
}
};
// u32
// test no data
try testing.expectEqual(
@as(?usize, 0),
upperBound(u32, @as(u32, 0), &[_]u32{}, {}, S.upper_u32),
);
// test below the first element
try testing.expectEqual(
@as(?usize, 0),
upperBound(u32, @as(u32, 0), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
);
// test equal to the first element
try testing.expectEqual(
@as(?usize, 1),
upperBound(u32, @as(u32, 2), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
);
// test between two numbers
try testing.expectEqual(
@as(?usize, 2),
upperBound(u32, @as(u32, 5), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
);
// test equal to a number not at the ends
try testing.expectEqual(
@as(?usize, 3),
upperBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
);
// test equal to the last element
try testing.expectEqual(
@as(?usize, 6),
upperBound(u32, @as(u32, 64), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
);
// test above the last element
try testing.expectEqual(
@as(?usize, 6),
upperBound(u32, @as(u32, 100), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
);
// i32
try testing.expectEqual(
@as(?usize, 2),
upperBound(i32, @as(i32, 5), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_i32),
);
// f32
try testing.expectEqual(
@as(?usize, 1),
upperBound(f32, @as(f32, -33.4), &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, {}, S.upper_f32),
);
}
/// Returns a range containing all elements equivalent to value in the range [first,last)
/// The function optimizes the number of comparisons performed by using a binary search O(2 * log n).
/// Example lessThan & greaterThan functions:
/// fn lower_u32(_: void, key: u32, lhs: u32) bool { return lhs < key; }
/// fn upper_u32(_: void, key: u32, rhs: u32) bool { return key >= rhs; }
pub fn equalRange(
comptime T: type,
key: anytype,
items: []const T,
context: anytype,
comptime lessThan: fn (context: @TypeOf(context), key: @TypeOf(key), lhs: T) bool,
comptime greaterThan: fn (context: @TypeOf(context), key: @TypeOf(key), rhs: T) bool,
) struct { usize, usize } {
return .{
lowerBound(T, key, items, context, lessThan),
upperBound(T, key, items, context, greaterThan),
};
}
test "equalRange" {
const S = struct {
fn lower_i32(context: void, key: i32, lhs: i32) bool {
_ = context;
return lhs < key;
}
fn upper_i32(context: void, key: i32, rhs: i32) bool {
_ = context;
return key >= rhs;
}
fn lower_u32(context: void, key: u32, lhs: u32) bool {
_ = context;
return lhs < key;
}
fn upper_u32(context: void, key: u32, rhs: u32) bool {
_ = context;
return key >= rhs;
}
fn lower_f32(context: void, key: f32, lhs: f32) bool {
_ = context;
return lhs < key;
}
fn upper_f32(context: void, key: f32, rhs: f32) bool {
_ = context;
return key >= rhs;
}
};
// i32
// test no data
try testing.expectEqual(
@as(struct { usize, usize }, .{ 0, 0 }),
equalRange(i32, @as(i32, 0), &[_]i32{}, {}, S.lower_i32, S.upper_i32),
);
// test below the first element
try testing.expectEqual(
@as(struct { usize, usize }, .{ 0, 0 }),
equalRange(i32, @as(i32, 0), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
);
// test equal to the first element
try testing.expectEqual(
@as(struct { usize, usize }, .{ 0, 1 }),
equalRange(i32, @as(i32, 2), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
);
// test between two numbers
try testing.expectEqual(
@as(struct { usize, usize }, .{ 2, 2 }),
equalRange(i32, @as(i32, 5), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
);
// test equal to a number not at the ends
try testing.expectEqual(
@as(struct { usize, usize }, .{ 2, 3 }),
equalRange(i32, @as(i32, 8), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
);
// test equal to the last element
try testing.expectEqual(
@as(struct { usize, usize }, .{ 5, 6 }),
equalRange(i32, @as(i32, 64), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
);
// test above the last element
try testing.expectEqual(
@as(struct { usize, usize }, .{ 6, 6 }),
equalRange(i32, @as(i32, 100), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
);
// test many of the same element
try testing.expectEqual(
@as(struct { usize, usize }, .{ 2, 6 }),
equalRange(i32, @as(i32, 8), &[_]i32{ 2, 4, 8, 8, 8, 8, 15, 22 }, {}, S.lower_i32, S.upper_i32),
);
// u32
try testing.expectEqual(
@as(struct { usize, usize }, .{ 2, 2 }),
equalRange(u32, @as(u32, 5), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32, S.upper_u32),
);
// f32
try testing.expectEqual(
@as(struct { usize, usize }, .{ 1, 1 }),
equalRange(f32, @as(f32, -33.4), &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, {}, S.lower_f32, S.upper_f32),
);
}
pub fn argMin(
comptime T: type,
items: []const T,