authorgravatar for craigoconnor@Craigs-MacBook-Pro.localCraig O'Connor <craigoconnor@Craigs-MacBook-Pro.local> 2024-01-27 22:42:56+01:00
committergravatar for john.schmidt.h@gmail.comJohn Schmidt <john.schmidt.h@gmail.com> 2024-02-07 21:00:24+01:00
log664c18544c2622352a4e38254496a349006d331b
tree16af022be1c0ac38ff5135183f759d5c25372110
parentee36131e6a8bf8611f8a6fd55116e98eae2ed63c

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.

1 files changed, 284 insertions(+), 0 deletions(-)

lib/std/sort.zig+284
...@@ -502,6 +502,290 @@ test "binarySearch" {...@@ -502,6 +502,290 @@ test "binarySearch" {
502 );502 );
503}503}
504504
505/// Returns the index pointing to the first element in the range [first,last)
506/// which does not compare less than val.
507/// The function optimizes the number of comparisons performed by using a binary search O(log n).
508/// An example lessThan function:
509/// fn lower_u32(_: void, key: u32, lhs: u32) bool { return lhs < key; }
510pub fn lowerBound(
511 comptime T: type,
512 key: anytype,
513 items: []const T,
514 context: anytype,
515 comptime lessThan: fn (context: @TypeOf(context), key: @TypeOf(key), lhs: T) bool,
516) usize {
517 var left: usize = 0;
518 var right: usize = items.len;
519 var mid: usize = undefined;
520
521 while (left < right) {
522 mid = left + (right - left) / 2;
523 if (lessThan(context, key, items[mid])) {
524 left = mid + 1;
525 } else {
526 right = mid;
527 }
528 }
529
530 return left;
531}
532
533test "lowerBound" {
534 const S = struct {
535 fn lower_u32(context: void, key: u32, lhs: u32) bool {
536 _ = context;
537 return lhs < key;
538 }
539 fn lower_i32(context: void, key: i32, lhs: i32) bool {
540 _ = context;
541 return lhs < key;
542 }
543 fn lower_f32(context: void, key: f32, lhs: f32) bool {
544 _ = context;
545 return lhs < key;
546 }
547 };
548
549 // u32
550 // test no data
551 try testing.expectEqual(
552 @as(?usize, 0),
553 lowerBound(u32, @as(u32, 0), &[_]u32{}, {}, S.lower_u32),
554 );
555 // test below the first element
556 try testing.expectEqual(
557 @as(?usize, 0),
558 lowerBound(u32, @as(u32, 0), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
559 );
560 // test equal to the first element
561 try testing.expectEqual(
562 @as(?usize, 0),
563 lowerBound(u32, @as(u32, 2), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
564 );
565 // test between two numbers
566 try testing.expectEqual(
567 @as(?usize, 2),
568 lowerBound(u32, @as(u32, 5), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
569 );
570 // test equal to a number not at the ends
571 try testing.expectEqual(
572 @as(?usize, 2),
573 lowerBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
574 );
575 // test equal to the last element
576 try testing.expectEqual(
577 @as(?usize, 5),
578 lowerBound(u32, @as(u32, 64), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
579 );
580 // test above the last element
581 try testing.expectEqual(
582 @as(?usize, 6),
583 lowerBound(u32, @as(u32, 100), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
584 );
585 // i32
586 try testing.expectEqual(
587 @as(?usize, 2),
588 lowerBound(i32, @as(i32, 5), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
589 );
590 // f32
591 try testing.expectEqual(
592 @as(?usize, 1),
593 lowerBound(f32, @as(f32, -33.4), &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, {}, S.lower_f32),
594 );
595}
596
597/// Returns the index pointing to the first element in the range [first,last)
598/// which compares greater than val.
599/// The function optimizes the number of comparisons performed by using a binary search O(log n).
600/// An example greaterThan function:
601/// fn upper_u32(_: void, key: u32, rhs: u32) bool { return key >= rhs; }
602pub fn upperBound(
603 comptime T: type,
604 key: anytype,
605 items: []const T,
606 context: anytype,
607 comptime greaterThan: fn (context: @TypeOf(context), key: @TypeOf(key), rhs: T) bool,
608) usize {
609 var left: usize = 0;
610 var right: usize = items.len;
611 var mid: usize = undefined;
612
613 while (left < right) {
614 mid = (right + left) / 2;
615 if (greaterThan(context, key, items[mid])) {
616 left = mid + 1;
617 } else {
618 right = mid;
619 }
620 }
621
622 return left;
623}
624
625test "upperBound" {
626 const S = struct {
627 fn upper_u32(context: void, key: u32, rhs: u32) bool {
628 _ = context;
629 return key >= rhs;
630 }
631 fn upper_i32(context: void, key: i32, rhs: i32) bool {
632 _ = context;
633 return key >= rhs;
634 }
635 fn upper_f32(context: void, key: f32, rhs: f32) bool {
636 _ = context;
637 return key >= rhs;
638 }
639 };
640
641 // u32
642 // test no data
643 try testing.expectEqual(
644 @as(?usize, 0),
645 upperBound(u32, @as(u32, 0), &[_]u32{}, {}, S.upper_u32),
646 );
647 // test below the first element
648 try testing.expectEqual(
649 @as(?usize, 0),
650 upperBound(u32, @as(u32, 0), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
651 );
652 // test equal to the first element
653 try testing.expectEqual(
654 @as(?usize, 1),
655 upperBound(u32, @as(u32, 2), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
656 );
657 // test between two numbers
658 try testing.expectEqual(
659 @as(?usize, 2),
660 upperBound(u32, @as(u32, 5), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
661 );
662 // test equal to a number not at the ends
663 try testing.expectEqual(
664 @as(?usize, 3),
665 upperBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
666 );
667 // test equal to the last element
668 try testing.expectEqual(
669 @as(?usize, 6),
670 upperBound(u32, @as(u32, 64), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
671 );
672 // test above the last element
673 try testing.expectEqual(
674 @as(?usize, 6),
675 upperBound(u32, @as(u32, 100), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_u32),
676 );
677 // i32
678 try testing.expectEqual(
679 @as(?usize, 2),
680 upperBound(i32, @as(i32, 5), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.upper_i32),
681 );
682 // f32
683 try testing.expectEqual(
684 @as(?usize, 1),
685 upperBound(f32, @as(f32, -33.4), &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, {}, S.upper_f32),
686 );
687}
688
689/// Returns a range containing all elements equivalent to value in the range [first,last)
690/// The function optimizes the number of comparisons performed by using a binary search O(2 * log n).
691/// Example lessThan & greaterThan functions:
692/// fn lower_u32(_: void, key: u32, lhs: u32) bool { return lhs < key; }
693/// fn upper_u32(_: void, key: u32, rhs: u32) bool { return key >= rhs; }
694pub fn equalRange(
695 comptime T: type,
696 key: anytype,
697 items: []const T,
698 context: anytype,
699 comptime lessThan: fn (context: @TypeOf(context), key: @TypeOf(key), lhs: T) bool,
700 comptime greaterThan: fn (context: @TypeOf(context), key: @TypeOf(key), rhs: T) bool,
701) struct { usize, usize } {
702 return .{
703 lowerBound(T, key, items, context, lessThan),
704 upperBound(T, key, items, context, greaterThan),
705 };
706}
707
708test "equalRange" {
709 const S = struct {
710 fn lower_i32(context: void, key: i32, lhs: i32) bool {
711 _ = context;
712 return lhs < key;
713 }
714 fn upper_i32(context: void, key: i32, rhs: i32) bool {
715 _ = context;
716 return key >= rhs;
717 }
718 fn lower_u32(context: void, key: u32, lhs: u32) bool {
719 _ = context;
720 return lhs < key;
721 }
722 fn upper_u32(context: void, key: u32, rhs: u32) bool {
723 _ = context;
724 return key >= rhs;
725 }
726 fn lower_f32(context: void, key: f32, lhs: f32) bool {
727 _ = context;
728 return lhs < key;
729 }
730 fn upper_f32(context: void, key: f32, rhs: f32) bool {
731 _ = context;
732 return key >= rhs;
733 }
734 };
735
736 // i32
737 // test no data
738 try testing.expectEqual(
739 @as(struct { usize, usize }, .{ 0, 0 }),
740 equalRange(i32, @as(i32, 0), &[_]i32{}, {}, S.lower_i32, S.upper_i32),
741 );
742 // test below the first element
743 try testing.expectEqual(
744 @as(struct { usize, usize }, .{ 0, 0 }),
745 equalRange(i32, @as(i32, 0), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
746 );
747 // test equal to the first element
748 try testing.expectEqual(
749 @as(struct { usize, usize }, .{ 0, 1 }),
750 equalRange(i32, @as(i32, 2), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
751 );
752 // test between two numbers
753 try testing.expectEqual(
754 @as(struct { usize, usize }, .{ 2, 2 }),
755 equalRange(i32, @as(i32, 5), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
756 );
757 // test equal to a number not at the ends
758 try testing.expectEqual(
759 @as(struct { usize, usize }, .{ 2, 3 }),
760 equalRange(i32, @as(i32, 8), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
761 );
762 // test equal to the last element
763 try testing.expectEqual(
764 @as(struct { usize, usize }, .{ 5, 6 }),
765 equalRange(i32, @as(i32, 64), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
766 );
767 // test above the last element
768 try testing.expectEqual(
769 @as(struct { usize, usize }, .{ 6, 6 }),
770 equalRange(i32, @as(i32, 100), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32, S.upper_i32),
771 );
772 // test many of the same element
773 try testing.expectEqual(
774 @as(struct { usize, usize }, .{ 2, 6 }),
775 equalRange(i32, @as(i32, 8), &[_]i32{ 2, 4, 8, 8, 8, 8, 15, 22 }, {}, S.lower_i32, S.upper_i32),
776 );
777 // u32
778 try testing.expectEqual(
779 @as(struct { usize, usize }, .{ 2, 2 }),
780 equalRange(u32, @as(u32, 5), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32, S.upper_u32),
781 );
782 // f32
783 try testing.expectEqual(
784 @as(struct { usize, usize }, .{ 1, 1 }),
785 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),
786 );
787}
788
505pub fn argMin(789pub fn argMin(
506 comptime T: type,790 comptime T: type,
507 items: []const T,791 items: []const T,