authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-02-07 18:48:41-08:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2024-02-07 18:48:41-08:00
logba8375328cd6ddf5025a400fc1d1817b88be6d53
tree26e3d7f135b236e734d91063834128f1547330b7
parent42fcca49c55eb9250a9ac8868fa6a9201dd9eb7e
parente487b576fadc2fd48100667ff233e4a0bb4f0aa5
signaturebadge-check Signed by PGP key B5690EEEBB952194

Merge pull request #18718 from schmee/bounds

Add upperBound, lowerBound, and equalRange

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

lib/std/sort.zig+268
...@@ -399,6 +399,13 @@ test "sort fuzz testing" {...@@ -399,6 +399,13 @@ test "sort fuzz testing" {
399 }399 }
400}400}
401401
402/// Returns the index of an element in `items` equal to `key`.
403/// If there are multiple such elements, returns the index of any one of them.
404/// If there are no such elements, returns `null`.
405///
406/// `items` must be sorted in ascending order with respect to `compareFn`.
407///
408/// O(log n) complexity.
402pub fn binarySearch(409pub fn binarySearch(
403 comptime T: type,410 comptime T: type,
404 key: anytype,411 key: anytype,
...@@ -502,6 +509,267 @@ test "binarySearch" {...@@ -502,6 +509,267 @@ test "binarySearch" {
502 );509 );
503}510}
504511
512/// Returns the index of the first element in `items` greater than or equal to `key`,
513/// or `items.len` if all elements are less than `key`.
514///
515/// `items` must be sorted in ascending order with respect to `compareFn`.
516///
517/// O(log n) complexity.
518pub fn lowerBound(
519 comptime T: type,
520 key: anytype,
521 items: []const T,
522 context: anytype,
523 comptime lessThan: fn (context: @TypeOf(context), lhs: @TypeOf(key), rhs: T) bool,
524) usize {
525 var left: usize = 0;
526 var right: usize = items.len;
527
528 while (left < right) {
529 const mid = left + (right - left) / 2;
530 if (lessThan(context, items[mid], key)) {
531 left = mid + 1;
532 } else {
533 right = mid;
534 }
535 }
536
537 return left;
538}
539
540test "lowerBound" {
541 const S = struct {
542 fn lower_u32(context: void, lhs: u32, rhs: u32) bool {
543 _ = context;
544 return lhs < rhs;
545 }
546 fn lower_i32(context: void, lhs: i32, rhs: i32) bool {
547 _ = context;
548 return lhs < rhs;
549 }
550 fn lower_f32(context: void, lhs: f32, rhs: f32) bool {
551 _ = context;
552 return lhs < rhs;
553 }
554 };
555
556 try testing.expectEqual(
557 @as(usize, 0),
558 lowerBound(u32, @as(u32, 0), &[_]u32{}, {}, S.lower_u32),
559 );
560 try testing.expectEqual(
561 @as(usize, 0),
562 lowerBound(u32, @as(u32, 0), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
563 );
564 try testing.expectEqual(
565 @as(usize, 0),
566 lowerBound(u32, @as(u32, 2), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
567 );
568 try testing.expectEqual(
569 @as(usize, 2),
570 lowerBound(u32, @as(u32, 5), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
571 );
572 try testing.expectEqual(
573 @as(usize, 2),
574 lowerBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
575 );
576 try testing.expectEqual(
577 @as(usize, 6),
578 lowerBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 7, 7, 7, 7, 16, 32, 64 }, {}, S.lower_u32),
579 );
580 try testing.expectEqual(
581 @as(usize, 2),
582 lowerBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 8, 8, 8, 8, 16, 32, 64 }, {}, S.lower_u32),
583 );
584 try testing.expectEqual(
585 @as(usize, 5),
586 lowerBound(u32, @as(u32, 64), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
587 );
588 try testing.expectEqual(
589 @as(usize, 6),
590 lowerBound(u32, @as(u32, 100), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
591 );
592 try testing.expectEqual(
593 @as(usize, 2),
594 lowerBound(i32, @as(i32, 5), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
595 );
596 try testing.expectEqual(
597 @as(usize, 1),
598 lowerBound(f32, @as(f32, -33.4), &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, {}, S.lower_f32),
599 );
600}
601
602/// Returns the index of the first element in `items` greater than `key`,
603/// or `items.len` if all elements are less than or equal to `key`.
604///
605/// `items` must be sorted in ascending order with respect to `compareFn`.
606///
607/// O(log n) complexity.
608pub fn upperBound(
609 comptime T: type,
610 key: anytype,
611 items: []const T,
612 context: anytype,
613 comptime lessThan: fn (context: @TypeOf(context), lhs: @TypeOf(key), rhs: T) bool,
614) usize {
615 var left: usize = 0;
616 var right: usize = items.len;
617
618 while (left < right) {
619 const mid = (right + left) / 2;
620 if (!lessThan(context, key, items[mid])) {
621 left = mid + 1;
622 } else {
623 right = mid;
624 }
625 }
626
627 return left;
628}
629
630test "upperBound" {
631 const S = struct {
632 fn lower_u32(context: void, lhs: u32, rhs: u32) bool {
633 _ = context;
634 return lhs < rhs;
635 }
636 fn lower_i32(context: void, lhs: i32, rhs: i32) bool {
637 _ = context;
638 return lhs < rhs;
639 }
640 fn lower_f32(context: void, lhs: f32, rhs: f32) bool {
641 _ = context;
642 return lhs < rhs;
643 }
644 };
645
646 try testing.expectEqual(
647 @as(usize, 0),
648 upperBound(u32, @as(u32, 0), &[_]u32{}, {}, S.lower_u32),
649 );
650 try testing.expectEqual(
651 @as(usize, 0),
652 upperBound(u32, @as(u32, 0), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
653 );
654 try testing.expectEqual(
655 @as(usize, 1),
656 upperBound(u32, @as(u32, 2), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
657 );
658 try testing.expectEqual(
659 @as(usize, 2),
660 upperBound(u32, @as(u32, 5), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
661 );
662 try testing.expectEqual(
663 @as(usize, 6),
664 upperBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 7, 7, 7, 7, 16, 32, 64 }, {}, S.lower_u32),
665 );
666 try testing.expectEqual(
667 @as(usize, 6),
668 upperBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 8, 8, 8, 8, 16, 32, 64 }, {}, S.lower_u32),
669 );
670 try testing.expectEqual(
671 @as(usize, 3),
672 upperBound(u32, @as(u32, 8), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
673 );
674 try testing.expectEqual(
675 @as(usize, 6),
676 upperBound(u32, @as(u32, 64), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
677 );
678 try testing.expectEqual(
679 @as(usize, 6),
680 upperBound(u32, @as(u32, 100), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
681 );
682 try testing.expectEqual(
683 @as(usize, 2),
684 upperBound(i32, @as(i32, 5), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
685 );
686 try testing.expectEqual(
687 @as(usize, 1),
688 upperBound(f32, @as(f32, -33.4), &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, {}, S.lower_f32),
689 );
690}
691
692/// Returns a tuple of the lower and upper indices in `items` between which all elements are equal to `key`.
693/// If no element in `items` is equal to `key`, both indices are the
694/// index of the first element in `items` greater than `key`.
695/// If no element in `items` is greater than `key`, both indices equal `items.len`.
696///
697/// `items` must be sorted in ascending order with respect to `compareFn`.
698///
699/// O(log n) complexity.
700///
701/// See also: `lowerBound` and `upperBound`.
702pub fn equalRange(
703 comptime T: type,
704 key: anytype,
705 items: []const T,
706 context: anytype,
707 comptime lessThan: fn (context: @TypeOf(context), lhs: @TypeOf(key), rhs: T) bool,
708) struct { usize, usize } {
709 return .{
710 lowerBound(T, key, items, context, lessThan),
711 upperBound(T, key, items, context, lessThan),
712 };
713}
714
715test "equalRange" {
716 const S = struct {
717 fn lower_u32(context: void, lhs: u32, rhs: u32) bool {
718 _ = context;
719 return lhs < rhs;
720 }
721 fn lower_i32(context: void, lhs: i32, rhs: i32) bool {
722 _ = context;
723 return lhs < rhs;
724 }
725 fn lower_f32(context: void, lhs: f32, rhs: f32) bool {
726 _ = context;
727 return lhs < rhs;
728 }
729 };
730
731 try testing.expectEqual(
732 @as(struct { usize, usize }, .{ 0, 0 }),
733 equalRange(i32, @as(i32, 0), &[_]i32{}, {}, S.lower_i32),
734 );
735 try testing.expectEqual(
736 @as(struct { usize, usize }, .{ 0, 0 }),
737 equalRange(i32, @as(i32, 0), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
738 );
739 try testing.expectEqual(
740 @as(struct { usize, usize }, .{ 0, 1 }),
741 equalRange(i32, @as(i32, 2), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
742 );
743 try testing.expectEqual(
744 @as(struct { usize, usize }, .{ 2, 2 }),
745 equalRange(i32, @as(i32, 5), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
746 );
747 try testing.expectEqual(
748 @as(struct { usize, usize }, .{ 2, 3 }),
749 equalRange(i32, @as(i32, 8), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
750 );
751 try testing.expectEqual(
752 @as(struct { usize, usize }, .{ 5, 6 }),
753 equalRange(i32, @as(i32, 64), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
754 );
755 try testing.expectEqual(
756 @as(struct { usize, usize }, .{ 6, 6 }),
757 equalRange(i32, @as(i32, 100), &[_]i32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_i32),
758 );
759 try testing.expectEqual(
760 @as(struct { usize, usize }, .{ 2, 6 }),
761 equalRange(i32, @as(i32, 8), &[_]i32{ 2, 4, 8, 8, 8, 8, 15, 22 }, {}, S.lower_i32),
762 );
763 try testing.expectEqual(
764 @as(struct { usize, usize }, .{ 2, 2 }),
765 equalRange(u32, @as(u32, 5), &[_]u32{ 2, 4, 8, 16, 32, 64 }, {}, S.lower_u32),
766 );
767 try testing.expectEqual(
768 @as(struct { usize, usize }, .{ 1, 1 }),
769 equalRange(f32, @as(f32, -33.4), &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, {}, S.lower_f32),
770 );
771}
772
505pub fn argMin(773pub fn argMin(
506 comptime T: type,774 comptime T: type,
507 items: []const T,775 items: []const T,