master
  1const std = @import("std.zig");
  2const assert = std.debug.assert;
  3const testing = std.testing;
  4const mem = std.mem;
  5const math = std.math;
  6
  7pub const Mode = enum { stable, unstable };
  8
  9pub const block = @import("sort/block.zig").block;
 10pub const pdq = @import("sort/pdq.zig").pdq;
 11pub const pdqContext = @import("sort/pdq.zig").pdqContext;
 12
 13/// Stable in-place sort. O(n) best case, O(pow(n, 2)) worst case.
 14/// O(1) memory (no allocator required).
 15/// Sorts in ascending order with respect to the given `lessThan` function.
 16pub fn insertion(
 17    comptime T: type,
 18    items: []T,
 19    context: anytype,
 20    comptime lessThanFn: fn (@TypeOf(context), lhs: T, rhs: T) bool,
 21) void {
 22    const Context = struct {
 23        items: []T,
 24        sub_ctx: @TypeOf(context),
 25
 26        pub fn lessThan(ctx: @This(), a: usize, b: usize) bool {
 27            return lessThanFn(ctx.sub_ctx, ctx.items[a], ctx.items[b]);
 28        }
 29
 30        pub fn swap(ctx: @This(), a: usize, b: usize) void {
 31            return mem.swap(T, &ctx.items[a], &ctx.items[b]);
 32        }
 33    };
 34    insertionContext(0, items.len, Context{ .items = items, .sub_ctx = context });
 35}
 36
 37/// Stable in-place sort. O(n) best case, O(pow(n, 2)) worst case.
 38/// O(1) memory (no allocator required).
 39/// `context` must have methods `swap` and `lessThan`,
 40/// which each take 2 `usize` parameters indicating the index of an item.
 41/// Sorts in ascending order with respect to `lessThan`.
 42pub fn insertionContext(a: usize, b: usize, context: anytype) void {
 43    assert(a <= b);
 44
 45    var i = a + 1;
 46    while (i < b) : (i += 1) {
 47        var j = i;
 48        while (j > a and context.lessThan(j, j - 1)) : (j -= 1) {
 49            context.swap(j, j - 1);
 50        }
 51    }
 52}
 53
 54/// Unstable in-place sort. O(n*log(n)) best case, worst case and average case.
 55/// O(1) memory (no allocator required).
 56/// Sorts in ascending order with respect to the given `lessThan` function.
 57pub fn heap(
 58    comptime T: type,
 59    items: []T,
 60    context: anytype,
 61    comptime lessThanFn: fn (@TypeOf(context), lhs: T, rhs: T) bool,
 62) void {
 63    const Context = struct {
 64        items: []T,
 65        sub_ctx: @TypeOf(context),
 66
 67        pub fn lessThan(ctx: @This(), a: usize, b: usize) bool {
 68            return lessThanFn(ctx.sub_ctx, ctx.items[a], ctx.items[b]);
 69        }
 70
 71        pub fn swap(ctx: @This(), a: usize, b: usize) void {
 72            return mem.swap(T, &ctx.items[a], &ctx.items[b]);
 73        }
 74    };
 75    heapContext(0, items.len, Context{ .items = items, .sub_ctx = context });
 76}
 77
 78/// Unstable in-place sort. O(n*log(n)) best case, worst case and average case.
 79/// O(1) memory (no allocator required).
 80/// `context` must have methods `swap` and `lessThan`,
 81/// which each take 2 `usize` parameters indicating the index of an item.
 82/// Sorts in ascending order with respect to `lessThan`.
 83pub fn heapContext(a: usize, b: usize, context: anytype) void {
 84    assert(a <= b);
 85    // build the heap in linear time.
 86    var i = a + (b - a) / 2;
 87    while (i > a) {
 88        i -= 1;
 89        siftDown(a, i, b, context);
 90    }
 91
 92    // pop maximal elements from the heap.
 93    i = b;
 94    while (i > a) {
 95        i -= 1;
 96        context.swap(a, i);
 97        siftDown(a, a, i, context);
 98    }
 99}
100
101fn siftDown(a: usize, target: usize, b: usize, context: anytype) void {
102    var cur = target;
103    while (true) {
104        // When we don't overflow from the multiply below, the following expression equals (2*cur) - (2*a) + a + 1
105        // The `+ a + 1` is safe because:
106        //  for `a > 0` then `2a >= a + 1`.
107        //  for `a = 0`, the expression equals `2*cur+1`. `2*cur` is an even number, therefore adding 1 is safe.
108        var child = (math.mul(usize, cur - a, 2) catch break) + a + 1;
109
110        // stop if we overshot the boundary
111        if (!(child < b)) break;
112
113        // `next_child` is at most `b`, therefore no overflow is possible
114        const next_child = child + 1;
115
116        // store the greater child in `child`
117        if (next_child < b and context.lessThan(child, next_child)) {
118            child = next_child;
119        }
120
121        // stop if the Heap invariant holds at `cur`.
122        if (context.lessThan(child, cur)) break;
123
124        // swap `cur` with the greater child,
125        // move one step down, and continue sifting.
126        context.swap(child, cur);
127        cur = child;
128    }
129}
130
131/// Use to generate a comparator function for a given type. e.g. `sort(u8, slice, {}, asc(u8))`.
132pub fn asc(comptime T: type) fn (void, T, T) bool {
133    return struct {
134        pub fn inner(_: void, a: T, b: T) bool {
135            return a < b;
136        }
137    }.inner;
138}
139
140/// Use to generate a comparator function for a given type. e.g. `sort(u8, slice, {}, desc(u8))`.
141pub fn desc(comptime T: type) fn (void, T, T) bool {
142    return struct {
143        pub fn inner(_: void, a: T, b: T) bool {
144            return a > b;
145        }
146    }.inner;
147}
148
149const asc_u8 = asc(u8);
150const asc_i32 = asc(i32);
151const desc_u8 = desc(u8);
152const desc_i32 = desc(i32);
153
154const sort_funcs = &[_]fn (comptime type, anytype, anytype, comptime anytype) void{
155    block,
156    pdq,
157    insertion,
158    heap,
159};
160
161const context_sort_funcs = &[_]fn (usize, usize, anytype) void{
162    // blockContext,
163    pdqContext,
164    insertionContext,
165    heapContext,
166};
167
168const IdAndValue = struct {
169    id: usize,
170    value: i32,
171
172    fn lessThan(context: void, a: IdAndValue, b: IdAndValue) bool {
173        _ = context;
174        return a.value < b.value;
175    }
176};
177
178test "stable sort" {
179    const expected = [_]IdAndValue{
180        IdAndValue{ .id = 0, .value = 0 },
181        IdAndValue{ .id = 1, .value = 0 },
182        IdAndValue{ .id = 2, .value = 0 },
183        IdAndValue{ .id = 0, .value = 1 },
184        IdAndValue{ .id = 1, .value = 1 },
185        IdAndValue{ .id = 2, .value = 1 },
186        IdAndValue{ .id = 0, .value = 2 },
187        IdAndValue{ .id = 1, .value = 2 },
188        IdAndValue{ .id = 2, .value = 2 },
189    };
190
191    var cases = [_][9]IdAndValue{
192        [_]IdAndValue{
193            IdAndValue{ .id = 0, .value = 0 },
194            IdAndValue{ .id = 0, .value = 1 },
195            IdAndValue{ .id = 0, .value = 2 },
196            IdAndValue{ .id = 1, .value = 0 },
197            IdAndValue{ .id = 1, .value = 1 },
198            IdAndValue{ .id = 1, .value = 2 },
199            IdAndValue{ .id = 2, .value = 0 },
200            IdAndValue{ .id = 2, .value = 1 },
201            IdAndValue{ .id = 2, .value = 2 },
202        },
203        [_]IdAndValue{
204            IdAndValue{ .id = 0, .value = 2 },
205            IdAndValue{ .id = 0, .value = 1 },
206            IdAndValue{ .id = 0, .value = 0 },
207            IdAndValue{ .id = 1, .value = 2 },
208            IdAndValue{ .id = 1, .value = 1 },
209            IdAndValue{ .id = 1, .value = 0 },
210            IdAndValue{ .id = 2, .value = 2 },
211            IdAndValue{ .id = 2, .value = 1 },
212            IdAndValue{ .id = 2, .value = 0 },
213        },
214    };
215
216    for (&cases) |*case| {
217        block(IdAndValue, (case.*)[0..], {}, IdAndValue.lessThan);
218        for (case.*, 0..) |item, i| {
219            try testing.expect(item.id == expected[i].id);
220            try testing.expect(item.value == expected[i].value);
221        }
222    }
223}
224
225test "stable sort fuzz testing" {
226    var prng = std.Random.DefaultPrng.init(std.testing.random_seed);
227    const random = prng.random();
228    const test_case_count = 10;
229
230    for (0..test_case_count) |_| {
231        const array_size = random.intRangeLessThan(usize, 0, 1000);
232        const array = try testing.allocator.alloc(IdAndValue, array_size);
233        defer testing.allocator.free(array);
234        // Value is a small random numbers to create collisions.
235        // Id is a  reverse index to make sure sorting function only uses provided `lessThan`.
236        for (array, 0..) |*item, index| {
237            item.* = .{
238                .value = random.intRangeLessThan(i32, 0, 100),
239                .id = array_size - index,
240            };
241        }
242        block(IdAndValue, array, {}, IdAndValue.lessThan);
243        if (array_size > 0) {
244            for (array[0 .. array_size - 1], array[1..]) |x, y| {
245                try testing.expect(x.value <= y.value);
246                if (x.value == y.value) {
247                    try testing.expect(x.id > y.id);
248                }
249            }
250        }
251    }
252}
253
254test "sort" {
255    const u8cases = [_][]const []const u8{
256        &[_][]const u8{
257            "",
258            "",
259        },
260        &[_][]const u8{
261            "a",
262            "a",
263        },
264        &[_][]const u8{
265            "az",
266            "az",
267        },
268        &[_][]const u8{
269            "za",
270            "az",
271        },
272        &[_][]const u8{
273            "asdf",
274            "adfs",
275        },
276        &[_][]const u8{
277            "one",
278            "eno",
279        },
280    };
281
282    const i32cases = [_][]const []const i32{
283        &[_][]const i32{
284            &[_]i32{},
285            &[_]i32{},
286        },
287        &[_][]const i32{
288            &[_]i32{1},
289            &[_]i32{1},
290        },
291        &[_][]const i32{
292            &[_]i32{ 0, 1 },
293            &[_]i32{ 0, 1 },
294        },
295        &[_][]const i32{
296            &[_]i32{ 1, 0 },
297            &[_]i32{ 0, 1 },
298        },
299        &[_][]const i32{
300            &[_]i32{ 1, -1, 0 },
301            &[_]i32{ -1, 0, 1 },
302        },
303        &[_][]const i32{
304            &[_]i32{ 2, 1, 3 },
305            &[_]i32{ 1, 2, 3 },
306        },
307        &[_][]const i32{
308            &[_]i32{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 55, 32, 39, 58, 21, 88, 43, 22, 59 },
309            &[_]i32{ 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 21, 22, 32, 39, 43, 55, 58, 59, 88 },
310        },
311    };
312
313    inline for (sort_funcs) |sortFn| {
314        for (u8cases) |case| {
315            var buf: [20]u8 = undefined;
316            const slice = buf[0..case[0].len];
317            @memcpy(slice, case[0]);
318            sortFn(u8, slice, {}, asc_u8);
319            try testing.expect(mem.eql(u8, slice, case[1]));
320        }
321
322        for (i32cases) |case| {
323            var buf: [20]i32 = undefined;
324            const slice = buf[0..case[0].len];
325            @memcpy(slice, case[0]);
326            sortFn(i32, slice, {}, asc_i32);
327            try testing.expect(mem.eql(i32, slice, case[1]));
328        }
329    }
330}
331
332test "sort descending" {
333    const rev_cases = [_][]const []const i32{
334        &[_][]const i32{
335            &[_]i32{},
336            &[_]i32{},
337        },
338        &[_][]const i32{
339            &[_]i32{1},
340            &[_]i32{1},
341        },
342        &[_][]const i32{
343            &[_]i32{ 0, 1 },
344            &[_]i32{ 1, 0 },
345        },
346        &[_][]const i32{
347            &[_]i32{ 1, 0 },
348            &[_]i32{ 1, 0 },
349        },
350        &[_][]const i32{
351            &[_]i32{ 1, -1, 0 },
352            &[_]i32{ 1, 0, -1 },
353        },
354        &[_][]const i32{
355            &[_]i32{ 2, 1, 3 },
356            &[_]i32{ 3, 2, 1 },
357        },
358    };
359
360    inline for (sort_funcs) |sortFn| {
361        for (rev_cases) |case| {
362            var buf: [8]i32 = undefined;
363            const slice = buf[0..case[0].len];
364            @memcpy(slice, case[0]);
365            sortFn(i32, slice, {}, desc_i32);
366            try testing.expect(mem.eql(i32, slice, case[1]));
367        }
368    }
369}
370
371test "sort with context in the middle of a slice" {
372    const Context = struct {
373        items: []i32,
374
375        pub fn lessThan(ctx: @This(), a: usize, b: usize) bool {
376            return ctx.items[a] < ctx.items[b];
377        }
378
379        pub fn swap(ctx: @This(), a: usize, b: usize) void {
380            return mem.swap(i32, &ctx.items[a], &ctx.items[b]);
381        }
382    };
383
384    const i32cases = [_][]const []const i32{
385        &[_][]const i32{
386            &[_]i32{ 0, 1, 8, 3, 6, 5, 4, 2, 9, 7, 10, 55, 32, 39, 58, 21, 88, 43, 22, 59 },
387            &[_]i32{ 50, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 21, 22, 32, 39, 43, 55, 58, 59, 88 },
388        },
389    };
390
391    const ranges = [_]struct { start: usize, end: usize }{
392        .{ .start = 10, .end = 20 },
393        .{ .start = 1, .end = 11 },
394        .{ .start = 3, .end = 7 },
395    };
396
397    inline for (context_sort_funcs) |sortFn| {
398        for (i32cases) |case| {
399            for (ranges) |range| {
400                var buf: [20]i32 = undefined;
401                const slice = buf[0..case[0].len];
402                @memcpy(slice, case[0]);
403                sortFn(range.start, range.end, Context{ .items = slice });
404                try testing.expectEqualSlices(i32, case[1][range.start..range.end], slice[range.start..range.end]);
405            }
406        }
407    }
408}
409
410test "sort fuzz testing" {
411    var prng = std.Random.DefaultPrng.init(std.testing.random_seed);
412    const random = prng.random();
413    const test_case_count = 10;
414
415    inline for (sort_funcs) |sortFn| {
416        for (0..test_case_count) |_| {
417            const array_size = random.intRangeLessThan(usize, 0, 1000);
418            const array = try testing.allocator.alloc(i32, array_size);
419            defer testing.allocator.free(array);
420            // populate with random data
421            for (array) |*item| {
422                item.* = random.intRangeLessThan(i32, 0, 100);
423            }
424            sortFn(i32, array, {}, asc_i32);
425            try testing.expect(isSorted(i32, array, {}, asc_i32));
426        }
427    }
428}
429
430/// Returns the index of an element in `items` returning `.eq` when given to `compareFn`.
431/// - If there are multiple such elements, returns the index of any one of them.
432/// - If there are no such elements, returns `null`.
433///
434/// `items` must be sorted in ascending order with respect to `compareFn`:
435/// ```
436/// [0]                                                   [len]
437/// ┌───┬───┬─/ /─┬───┬───┬───┬─/ /─┬───┬───┬───┬─/ /─┬───┐
438/// │.lt│.lt│ \ \ │.lt│.eq│.eq│ \ \ │.eq│.gt│.gt│ \ \ │.gt│
439/// └───┴───┴─/ /─┴───┴───┴───┴─/ /─┴───┴───┴───┴─/ /─┴───┘
440/// ├─────────────────┼─────────────────┼─────────────────┤
441///  ↳ zero or more    ↳ zero or more    ↳ zero or more
442///                   ├─────────────────┤
443///                    ↳ if not null, returned
444///                      index is in this range
445/// ```
446///
447/// `O(log n)` time complexity.
448///
449/// See also: `lowerBound, `upperBound`, `partitionPoint`, `equalRange`.
450pub fn binarySearch(
451    comptime T: type,
452    items: []const T,
453    context: anytype,
454    comptime compareFn: fn (@TypeOf(context), T) std.math.Order,
455) ?usize {
456    var low: usize = 0;
457    var high: usize = items.len;
458
459    while (low < high) {
460        // Avoid overflowing in the midpoint calculation
461        const mid = low + (high - low) / 2;
462        switch (compareFn(context, items[mid])) {
463            .eq => return mid,
464            .gt => low = mid + 1,
465            .lt => high = mid,
466        }
467    }
468    return null;
469}
470
471test binarySearch {
472    const S = struct {
473        fn orderU32(context: u32, item: u32) std.math.Order {
474            return std.math.order(context, item);
475        }
476        fn orderI32(context: i32, item: i32) std.math.Order {
477            return std.math.order(context, item);
478        }
479        fn orderLength(context: usize, item: []const u8) std.math.Order {
480            return std.math.order(context, item.len);
481        }
482    };
483    const R = struct {
484        b: i32,
485        e: i32,
486
487        fn r(b: i32, e: i32) @This() {
488            return .{ .b = b, .e = e };
489        }
490
491        fn order(context: i32, item: @This()) std.math.Order {
492            if (context < item.b) {
493                return .lt;
494            } else if (context > item.e) {
495                return .gt;
496            } else {
497                return .eq;
498            }
499        }
500    };
501
502    try std.testing.expectEqual(null, binarySearch(u32, &[_]u32{}, @as(u32, 1), S.orderU32));
503    try std.testing.expectEqual(0, binarySearch(u32, &[_]u32{1}, @as(u32, 1), S.orderU32));
504    try std.testing.expectEqual(null, binarySearch(u32, &[_]u32{0}, @as(u32, 1), S.orderU32));
505    try std.testing.expectEqual(null, binarySearch(u32, &[_]u32{1}, @as(u32, 0), S.orderU32));
506    try std.testing.expectEqual(4, binarySearch(u32, &[_]u32{ 1, 2, 3, 4, 5 }, @as(u32, 5), S.orderU32));
507    try std.testing.expectEqual(0, binarySearch(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 2), S.orderU32));
508    try std.testing.expectEqual(1, binarySearch(i32, &[_]i32{ -7, -4, 0, 9, 10 }, @as(i32, -4), S.orderI32));
509    try std.testing.expectEqual(3, binarySearch(i32, &[_]i32{ -100, -25, 2, 98, 99, 100 }, @as(i32, 98), S.orderI32));
510    try std.testing.expectEqual(null, binarySearch(R, &[_]R{ R.r(-100, -50), R.r(-40, -20), R.r(-10, 20), R.r(30, 40) }, @as(i32, -45), R.order));
511    try std.testing.expectEqual(2, binarySearch(R, &[_]R{ R.r(-100, -50), R.r(-40, -20), R.r(-10, 20), R.r(30, 40) }, @as(i32, 10), R.order));
512    try std.testing.expectEqual(1, binarySearch(R, &[_]R{ R.r(-100, -50), R.r(-40, -20), R.r(-10, 20), R.r(30, 40) }, @as(i32, -20), R.order));
513    try std.testing.expectEqual(2, binarySearch([]const u8, &[_][]const u8{ "", "abc", "1234", "vwxyz" }, @as(usize, 4), S.orderLength));
514}
515
516/// Returns the index of the first element in `items` that is greater than or equal to `context`,
517/// as determined by `compareFn`. If no such element exists, returns `items.len`.
518///
519/// `items` must be sorted in ascending order with respect to `compareFn`:
520/// ```
521/// [0]                                                   [len]
522/// ┌───┬───┬─/ /─┬───┬───┬───┬─/ /─┬───┬───┬───┬─/ /─┬───┐
523/// │.lt│.lt│ \ \ │.lt│.eq│.eq│ \ \ │.eq│.gt│.gt│ \ \ │.gt│
524/// └───┴───┴─/ /─┴───┴───┴───┴─/ /─┴───┴───┴───┴─/ /─┴───┘
525/// ├─────────────────┼─────────────────┼─────────────────┤
526///  ↳ zero or more    ↳ zero or more    ↳ zero or more
527///                   ├───┤
528///                    ↳ returned index
529/// ```
530///
531/// `O(log n)` time complexity.
532///
533/// See also: `binarySearch`, `upperBound`, `partitionPoint`, `equalRange`.
534pub fn lowerBound(
535    comptime T: type,
536    items: []const T,
537    context: anytype,
538    comptime compareFn: fn (@TypeOf(context), T) std.math.Order,
539) usize {
540    const S = struct {
541        fn predicate(ctx: @TypeOf(context), item: T) bool {
542            return compareFn(ctx, item).invert() == .lt;
543        }
544    };
545    return partitionPoint(T, items, context, S.predicate);
546}
547
548test lowerBound {
549    const S = struct {
550        fn compareU32(context: u32, item: u32) std.math.Order {
551            return std.math.order(context, item);
552        }
553        fn compareI32(context: i32, item: i32) std.math.Order {
554            return std.math.order(context, item);
555        }
556        fn compareF32(context: f32, item: f32) std.math.Order {
557            return std.math.order(context, item);
558        }
559    };
560    const R = struct {
561        val: i32,
562
563        fn r(val: i32) @This() {
564            return .{ .val = val };
565        }
566
567        fn compareFn(context: i32, item: @This()) std.math.Order {
568            return std.math.order(context, item.val);
569        }
570    };
571
572    try std.testing.expectEqual(0, lowerBound(u32, &[_]u32{}, @as(u32, 0), S.compareU32));
573    try std.testing.expectEqual(0, lowerBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 0), S.compareU32));
574    try std.testing.expectEqual(0, lowerBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 2), S.compareU32));
575    try std.testing.expectEqual(2, lowerBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 5), S.compareU32));
576    try std.testing.expectEqual(2, lowerBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 8), S.compareU32));
577    try std.testing.expectEqual(6, lowerBound(u32, &[_]u32{ 2, 4, 7, 7, 7, 7, 16, 32, 64 }, @as(u32, 8), S.compareU32));
578    try std.testing.expectEqual(2, lowerBound(u32, &[_]u32{ 2, 4, 8, 8, 8, 8, 16, 32, 64 }, @as(u32, 8), S.compareU32));
579    try std.testing.expectEqual(5, lowerBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 64), S.compareU32));
580    try std.testing.expectEqual(6, lowerBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 100), S.compareU32));
581    try std.testing.expectEqual(2, lowerBound(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 5), S.compareI32));
582    try std.testing.expectEqual(1, lowerBound(f32, &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, @as(f32, -33.4), S.compareF32));
583    try std.testing.expectEqual(2, lowerBound(R, &[_]R{ R.r(-100), R.r(-40), R.r(-10), R.r(30) }, @as(i32, -20), R.compareFn));
584}
585
586/// Returns the index of the first element in `items` that is greater than `context`, as determined
587/// by `compareFn`. If no such element exists, returns `items.len`.
588///
589/// `items` must be sorted in ascending order with respect to `compareFn`:
590/// ```
591/// [0]                                                   [len]
592/// ┌───┬───┬─/ /─┬───┬───┬───┬─/ /─┬───┬───┬───┬─/ /─┬───┐
593/// │.lt│.lt│ \ \ │.lt│.eq│.eq│ \ \ │.eq│.gt│.gt│ \ \ │.gt│
594/// └───┴───┴─/ /─┴───┴───┴───┴─/ /─┴───┴───┴───┴─/ /─┴───┘
595/// ├─────────────────┼─────────────────┼─────────────────┤
596///  ↳ zero or more    ↳ zero or more    ↳ zero or more
597///                                     ├───┤
598///                                      ↳ returned index
599/// ```
600///
601/// `O(log n)` time complexity.
602///
603/// See also: `binarySearch`, `lowerBound`, `partitionPoint`, `equalRange`.
604pub fn upperBound(
605    comptime T: type,
606    items: []const T,
607    context: anytype,
608    comptime compareFn: fn (@TypeOf(context), T) std.math.Order,
609) usize {
610    const S = struct {
611        fn predicate(ctx: @TypeOf(context), item: T) bool {
612            return compareFn(ctx, item).invert() != .gt;
613        }
614    };
615    return partitionPoint(T, items, context, S.predicate);
616}
617
618test upperBound {
619    const S = struct {
620        fn compareU32(context: u32, item: u32) std.math.Order {
621            return std.math.order(context, item);
622        }
623        fn compareI32(context: i32, item: i32) std.math.Order {
624            return std.math.order(context, item);
625        }
626        fn compareF32(context: f32, item: f32) std.math.Order {
627            return std.math.order(context, item);
628        }
629    };
630    const R = struct {
631        val: i32,
632
633        fn r(val: i32) @This() {
634            return .{ .val = val };
635        }
636
637        fn compareFn(context: i32, item: @This()) std.math.Order {
638            return std.math.order(context, item.val);
639        }
640    };
641
642    try std.testing.expectEqual(0, upperBound(u32, &[_]u32{}, @as(u32, 0), S.compareU32));
643    try std.testing.expectEqual(0, upperBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 0), S.compareU32));
644    try std.testing.expectEqual(1, upperBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 2), S.compareU32));
645    try std.testing.expectEqual(2, upperBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 5), S.compareU32));
646    try std.testing.expectEqual(6, upperBound(u32, &[_]u32{ 2, 4, 7, 7, 7, 7, 16, 32, 64 }, @as(u32, 8), S.compareU32));
647    try std.testing.expectEqual(6, upperBound(u32, &[_]u32{ 2, 4, 8, 8, 8, 8, 16, 32, 64 }, @as(u32, 8), S.compareU32));
648    try std.testing.expectEqual(3, upperBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 8), S.compareU32));
649    try std.testing.expectEqual(6, upperBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 64), S.compareU32));
650    try std.testing.expectEqual(6, upperBound(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 100), S.compareU32));
651    try std.testing.expectEqual(2, upperBound(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 5), S.compareI32));
652    try std.testing.expectEqual(1, upperBound(f32, &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, @as(f32, -33.4), S.compareF32));
653    try std.testing.expectEqual(2, upperBound(R, &[_]R{ R.r(-100), R.r(-40), R.r(-10), R.r(30) }, @as(i32, -20), R.compareFn));
654}
655
656/// Returns the index of the partition point of `items` in relation to the given predicate.
657/// - If all elements of `items` satisfy the predicate the returned value is `items.len`.
658///
659/// `items` must contain a prefix for which all elements satisfy the predicate,
660/// and beyond which none of the elements satisfy the predicate:
661/// ```
662/// [0]                                          [len]
663/// ┌────┬────┬─/ /─┬────┬─────┬─────┬─/ /─┬─────┐
664/// │true│true│ \ \ │true│false│false│ \ \ │false│
665/// └────┴────┴─/ /─┴────┴─────┴─────┴─/ /─┴─────┘
666/// ├────────────────────┼───────────────────────┤
667///  ↳ zero or more       ↳ zero or more
668///                      ├─────┤
669///                       ↳ returned index
670/// ```
671///
672/// `O(log n)` time complexity.
673///
674/// See also: `binarySearch`, `lowerBound, `upperBound`, `equalRange`.
675pub fn partitionPoint(
676    comptime T: type,
677    items: []const T,
678    context: anytype,
679    comptime predicate: fn (@TypeOf(context), T) bool,
680) usize {
681    var it: usize = 0;
682    var len: usize = items.len;
683
684    while (len > 1) {
685        const half: usize = len / 2;
686        len -= half;
687        if (predicate(context, items[it + half - 1])) {
688            @branchHint(.unpredictable);
689            it += half;
690        }
691    }
692
693    if (it < items.len) {
694        it += @intFromBool(predicate(context, items[it]));
695    }
696
697    return it;
698}
699
700test partitionPoint {
701    const S = struct {
702        fn lowerU32(context: u32, item: u32) bool {
703            return item < context;
704        }
705        fn lowerI32(context: i32, item: i32) bool {
706            return item < context;
707        }
708        fn lowerF32(context: f32, item: f32) bool {
709            return item < context;
710        }
711        fn lowerEqU32(context: u32, item: u32) bool {
712            return item <= context;
713        }
714        fn lowerEqI32(context: i32, item: i32) bool {
715            return item <= context;
716        }
717        fn lowerEqF32(context: f32, item: f32) bool {
718            return item <= context;
719        }
720        fn isEven(_: void, item: u8) bool {
721            return item % 2 == 0;
722        }
723    };
724
725    try std.testing.expectEqual(0, partitionPoint(u32, &[_]u32{}, @as(u32, 0), S.lowerU32));
726    try std.testing.expectEqual(0, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 0), S.lowerU32));
727    try std.testing.expectEqual(0, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 2), S.lowerU32));
728    try std.testing.expectEqual(2, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 5), S.lowerU32));
729    try std.testing.expectEqual(2, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 8), S.lowerU32));
730    try std.testing.expectEqual(6, partitionPoint(u32, &[_]u32{ 2, 4, 7, 7, 7, 7, 16, 32, 64 }, @as(u32, 8), S.lowerU32));
731    try std.testing.expectEqual(2, partitionPoint(u32, &[_]u32{ 2, 4, 8, 8, 8, 8, 16, 32, 64 }, @as(u32, 8), S.lowerU32));
732    try std.testing.expectEqual(5, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 64), S.lowerU32));
733    try std.testing.expectEqual(6, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 100), S.lowerU32));
734    try std.testing.expectEqual(2, partitionPoint(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 5), S.lowerI32));
735    try std.testing.expectEqual(1, partitionPoint(f32, &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, @as(f32, -33.4), S.lowerF32));
736    try std.testing.expectEqual(0, partitionPoint(u32, &[_]u32{}, @as(u32, 0), S.lowerEqU32));
737    try std.testing.expectEqual(0, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 0), S.lowerEqU32));
738    try std.testing.expectEqual(1, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 2), S.lowerEqU32));
739    try std.testing.expectEqual(2, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 5), S.lowerEqU32));
740    try std.testing.expectEqual(6, partitionPoint(u32, &[_]u32{ 2, 4, 7, 7, 7, 7, 16, 32, 64 }, @as(u32, 8), S.lowerEqU32));
741    try std.testing.expectEqual(6, partitionPoint(u32, &[_]u32{ 2, 4, 8, 8, 8, 8, 16, 32, 64 }, @as(u32, 8), S.lowerEqU32));
742    try std.testing.expectEqual(3, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 8), S.lowerEqU32));
743    try std.testing.expectEqual(6, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 64), S.lowerEqU32));
744    try std.testing.expectEqual(6, partitionPoint(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 100), S.lowerEqU32));
745    try std.testing.expectEqual(2, partitionPoint(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 5), S.lowerEqI32));
746    try std.testing.expectEqual(1, partitionPoint(f32, &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, @as(f32, -33.4), S.lowerEqF32));
747    try std.testing.expectEqual(4, partitionPoint(u8, &[_]u8{ 0, 50, 14, 2, 5, 71 }, {}, S.isEven));
748}
749
750/// Returns a tuple of the lower and upper indices in `items` between which all
751/// elements return `.eq` when given to `compareFn`.
752/// - If no element in `items` returns `.eq`, both indices are the
753/// index of the first element in `items` returning `.gt`.
754/// - If no element in `items` returns `.gt`, both indices equal `items.len`.
755///
756/// `items` must be sorted in ascending order with respect to `compareFn`:
757/// ```
758/// [0]                                                   [len]
759/// ┌───┬───┬─/ /─┬───┬───┬───┬─/ /─┬───┬───┬───┬─/ /─┬───┐
760/// │.lt│.lt│ \ \ │.lt│.eq│.eq│ \ \ │.eq│.gt│.gt│ \ \ │.gt│
761/// └───┴───┴─/ /─┴───┴───┴───┴─/ /─┴───┴───┴───┴─/ /─┴───┘
762/// ├─────────────────┼─────────────────┼─────────────────┤
763///  ↳ zero or more    ↳ zero or more    ↳ zero or more
764///                   ├─────────────────┤
765///                    ↳ returned range
766/// ```
767///
768/// `O(log n)` time complexity.
769///
770/// See also: `binarySearch`, `lowerBound, `upperBound`, `partitionPoint`.
771pub fn equalRange(
772    comptime T: type,
773    items: []const T,
774    context: anytype,
775    comptime compareFn: fn (@TypeOf(context), T) std.math.Order,
776) struct { usize, usize } {
777    var low: usize = 0;
778    var high: usize = items.len;
779
780    while (low < high) {
781        const mid = low + (high - low) / 2;
782        switch (compareFn(context, items[mid])) {
783            .gt => {
784                low = mid + 1;
785            },
786            .lt => {
787                high = mid;
788            },
789            .eq => {
790                return .{
791                    low + std.sort.lowerBound(
792                        T,
793                        items[low..mid],
794                        context,
795                        compareFn,
796                    ),
797                    mid + std.sort.upperBound(
798                        T,
799                        items[mid..high],
800                        context,
801                        compareFn,
802                    ),
803                };
804            },
805        }
806    }
807
808    return .{ low, low };
809}
810
811test equalRange {
812    const S = struct {
813        fn orderU32(context: u32, item: u32) std.math.Order {
814            return std.math.order(context, item);
815        }
816        fn orderI32(context: i32, item: i32) std.math.Order {
817            return std.math.order(context, item);
818        }
819        fn orderF32(context: f32, item: f32) std.math.Order {
820            return std.math.order(context, item);
821        }
822        fn orderLength(context: usize, item: []const u8) std.math.Order {
823            return std.math.order(context, item.len);
824        }
825    };
826
827    try std.testing.expectEqual(.{ 0, 0 }, equalRange(i32, &[_]i32{}, @as(i32, 0), S.orderI32));
828    try std.testing.expectEqual(.{ 0, 0 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 0), S.orderI32));
829    try std.testing.expectEqual(.{ 0, 1 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 2), S.orderI32));
830    try std.testing.expectEqual(.{ 2, 2 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 5), S.orderI32));
831    try std.testing.expectEqual(.{ 2, 3 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 8), S.orderI32));
832    try std.testing.expectEqual(.{ 5, 6 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 64), S.orderI32));
833    try std.testing.expectEqual(.{ 6, 6 }, equalRange(i32, &[_]i32{ 2, 4, 8, 16, 32, 64 }, @as(i32, 100), S.orderI32));
834    try std.testing.expectEqual(.{ 2, 6 }, equalRange(i32, &[_]i32{ 2, 4, 8, 8, 8, 8, 15, 22 }, @as(i32, 8), S.orderI32));
835    try std.testing.expectEqual(.{ 2, 2 }, equalRange(u32, &[_]u32{ 2, 4, 8, 16, 32, 64 }, @as(u32, 5), S.orderU32));
836    try std.testing.expectEqual(.{ 3, 5 }, equalRange(u32, &[_]u32{ 2, 3, 4, 5, 5 }, @as(u32, 5), S.orderU32));
837    try std.testing.expectEqual(.{ 1, 1 }, equalRange(f32, &[_]f32{ -54.2, -26.7, 0.0, 56.55, 100.1, 322.0 }, @as(f32, -33.4), S.orderF32));
838    try std.testing.expectEqual(.{ 3, 5 }, equalRange(
839        []const u8,
840        &[_][]const u8{ "Mars", "Venus", "Earth", "Saturn", "Uranus", "Mercury", "Jupiter", "Neptune" },
841        @as(usize, 6),
842        S.orderLength,
843    ));
844}
845
846pub fn argMin(
847    comptime T: type,
848    items: []const T,
849    context: anytype,
850    comptime lessThan: fn (@TypeOf(context), lhs: T, rhs: T) bool,
851) ?usize {
852    if (items.len == 0) {
853        return null;
854    }
855
856    var smallest = items[0];
857    var smallest_index: usize = 0;
858    for (items[1..], 0..) |item, i| {
859        if (lessThan(context, item, smallest)) {
860            smallest = item;
861            smallest_index = i + 1;
862        }
863    }
864
865    return smallest_index;
866}
867
868test argMin {
869    try testing.expectEqual(@as(?usize, null), argMin(i32, &[_]i32{}, {}, asc_i32));
870    try testing.expectEqual(@as(?usize, 0), argMin(i32, &[_]i32{1}, {}, asc_i32));
871    try testing.expectEqual(@as(?usize, 0), argMin(i32, &[_]i32{ 1, 2, 3, 4, 5 }, {}, asc_i32));
872    try testing.expectEqual(@as(?usize, 3), argMin(i32, &[_]i32{ 9, 3, 8, 2, 5 }, {}, asc_i32));
873    try testing.expectEqual(@as(?usize, 0), argMin(i32, &[_]i32{ 1, 1, 1, 1, 1 }, {}, asc_i32));
874    try testing.expectEqual(@as(?usize, 0), argMin(i32, &[_]i32{ -10, 1, 10 }, {}, asc_i32));
875    try testing.expectEqual(@as(?usize, 3), argMin(i32, &[_]i32{ 6, 3, 5, 7, 6 }, {}, desc_i32));
876}
877
878pub fn min(
879    comptime T: type,
880    items: []const T,
881    context: anytype,
882    comptime lessThan: fn (context: @TypeOf(context), lhs: T, rhs: T) bool,
883) ?T {
884    const i = argMin(T, items, context, lessThan) orelse return null;
885    return items[i];
886}
887
888test min {
889    try testing.expectEqual(@as(?i32, null), min(i32, &[_]i32{}, {}, asc_i32));
890    try testing.expectEqual(@as(?i32, 1), min(i32, &[_]i32{1}, {}, asc_i32));
891    try testing.expectEqual(@as(?i32, 1), min(i32, &[_]i32{ 1, 2, 3, 4, 5 }, {}, asc_i32));
892    try testing.expectEqual(@as(?i32, 2), min(i32, &[_]i32{ 9, 3, 8, 2, 5 }, {}, asc_i32));
893    try testing.expectEqual(@as(?i32, 1), min(i32, &[_]i32{ 1, 1, 1, 1, 1 }, {}, asc_i32));
894    try testing.expectEqual(@as(?i32, -10), min(i32, &[_]i32{ -10, 1, 10 }, {}, asc_i32));
895    try testing.expectEqual(@as(?i32, 7), min(i32, &[_]i32{ 6, 3, 5, 7, 6 }, {}, desc_i32));
896}
897
898pub fn argMax(
899    comptime T: type,
900    items: []const T,
901    context: anytype,
902    comptime lessThan: fn (context: @TypeOf(context), lhs: T, rhs: T) bool,
903) ?usize {
904    if (items.len == 0) {
905        return null;
906    }
907
908    var biggest = items[0];
909    var biggest_index: usize = 0;
910    for (items[1..], 0..) |item, i| {
911        if (lessThan(context, biggest, item)) {
912            biggest = item;
913            biggest_index = i + 1;
914        }
915    }
916
917    return biggest_index;
918}
919
920test argMax {
921    try testing.expectEqual(@as(?usize, null), argMax(i32, &[_]i32{}, {}, asc_i32));
922    try testing.expectEqual(@as(?usize, 0), argMax(i32, &[_]i32{1}, {}, asc_i32));
923    try testing.expectEqual(@as(?usize, 4), argMax(i32, &[_]i32{ 1, 2, 3, 4, 5 }, {}, asc_i32));
924    try testing.expectEqual(@as(?usize, 0), argMax(i32, &[_]i32{ 9, 3, 8, 2, 5 }, {}, asc_i32));
925    try testing.expectEqual(@as(?usize, 0), argMax(i32, &[_]i32{ 1, 1, 1, 1, 1 }, {}, asc_i32));
926    try testing.expectEqual(@as(?usize, 2), argMax(i32, &[_]i32{ -10, 1, 10 }, {}, asc_i32));
927    try testing.expectEqual(@as(?usize, 1), argMax(i32, &[_]i32{ 6, 3, 5, 7, 6 }, {}, desc_i32));
928}
929
930pub fn max(
931    comptime T: type,
932    items: []const T,
933    context: anytype,
934    comptime lessThan: fn (context: @TypeOf(context), lhs: T, rhs: T) bool,
935) ?T {
936    const i = argMax(T, items, context, lessThan) orelse return null;
937    return items[i];
938}
939
940test max {
941    try testing.expectEqual(@as(?i32, null), max(i32, &[_]i32{}, {}, asc_i32));
942    try testing.expectEqual(@as(?i32, 1), max(i32, &[_]i32{1}, {}, asc_i32));
943    try testing.expectEqual(@as(?i32, 5), max(i32, &[_]i32{ 1, 2, 3, 4, 5 }, {}, asc_i32));
944    try testing.expectEqual(@as(?i32, 9), max(i32, &[_]i32{ 9, 3, 8, 2, 5 }, {}, asc_i32));
945    try testing.expectEqual(@as(?i32, 1), max(i32, &[_]i32{ 1, 1, 1, 1, 1 }, {}, asc_i32));
946    try testing.expectEqual(@as(?i32, 10), max(i32, &[_]i32{ -10, 1, 10 }, {}, asc_i32));
947    try testing.expectEqual(@as(?i32, 3), max(i32, &[_]i32{ 6, 3, 5, 7, 6 }, {}, desc_i32));
948}
949
950pub fn isSorted(
951    comptime T: type,
952    items: []const T,
953    context: anytype,
954    comptime lessThan: fn (context: @TypeOf(context), lhs: T, rhs: T) bool,
955) bool {
956    var i: usize = 1;
957    while (i < items.len) : (i += 1) {
958        if (lessThan(context, items[i], items[i - 1])) {
959            return false;
960        }
961    }
962
963    return true;
964}
965
966test isSorted {
967    try testing.expect(isSorted(i32, &[_]i32{}, {}, asc_i32));
968    try testing.expect(isSorted(i32, &[_]i32{10}, {}, asc_i32));
969    try testing.expect(isSorted(i32, &[_]i32{ 1, 2, 3, 4, 5 }, {}, asc_i32));
970    try testing.expect(isSorted(i32, &[_]i32{ -10, 1, 1, 1, 10 }, {}, asc_i32));
971
972    try testing.expect(isSorted(i32, &[_]i32{}, {}, desc_i32));
973    try testing.expect(isSorted(i32, &[_]i32{-20}, {}, desc_i32));
974    try testing.expect(isSorted(i32, &[_]i32{ 3, 2, 1, 0, -1 }, {}, desc_i32));
975    try testing.expect(isSorted(i32, &[_]i32{ 10, -10 }, {}, desc_i32));
976
977    try testing.expect(isSorted(i32, &[_]i32{ 1, 1, 1, 1, 1 }, {}, asc_i32));
978    try testing.expect(isSorted(i32, &[_]i32{ 1, 1, 1, 1, 1 }, {}, desc_i32));
979
980    try testing.expectEqual(false, isSorted(i32, &[_]i32{ 5, 4, 3, 2, 1 }, {}, asc_i32));
981    try testing.expectEqual(false, isSorted(i32, &[_]i32{ 1, 2, 3, 4, 5 }, {}, desc_i32));
982
983    try testing.expect(isSorted(u8, "abcd", {}, asc_u8));
984    try testing.expect(isSorted(u8, "zyxw", {}, desc_u8));
985
986    try testing.expectEqual(false, isSorted(u8, "abcd", {}, desc_u8));
987    try testing.expectEqual(false, isSorted(u8, "zyxw", {}, asc_u8));
988
989    try testing.expect(isSorted(u8, "ffff", {}, asc_u8));
990    try testing.expect(isSorted(u8, "ffff", {}, desc_u8));
991}