Building on @vadian's and @Martin R's answers, I noticed some minor discrepancies, mostly with the insertion index either not matching the index of an equivalent element in the collection, or of it not matching the first index of a sequence of equivalent elements.
For instance:
- If you wanted to find an insertion index for
5
in [4, 5, 6]
, the index 2
would be returned, which may be problematic if you want to simply search for the value.
- In
[5, 5, 5]
, searching once again for 5
returns the index 1
, which is not the first insertion index.
This doesn't match with the behavior of NSArray's implementation and its various options, so here is yet another solution that tries to take this into account:
extension RandomAccessCollection {
/// Get the index of or an insertion index for a new element in
/// a sorted collection in ascending order.
/// - Parameter value: The element to insert into the array.
/// - Returns: The index suitable for inserting the new element
/// into the array, or the first index of an existing element.
@inlinable
public func sortedInsertionIndex(
of element: Element
) -> Index where Element: Comparable {
sortedInsertionIndex(of: element, by: <)
}
/// Get the index of or an insertion index for a new element in
/// a sorted collection that matches the rule defined by the predicate.
/// - Parameters:
/// - value: The element to insert into the array.
/// - areInIncreasingOrder:
/// A closure that determines if the first element should
/// come before the second element. For instance: `<`.
/// - Returns: The index suitable for inserting the new element
/// into the array, or the first index of an existing element.
@inlinable
public func sortedInsertionIndex(
of element: Element,
by areInIncreasingOrder: (Element, Element) throws -> Bool
) rethrows -> Index {
try sortedInsertionIndex { try areInIncreasingOrder($0, element) }
}
/// Get the index of or an insertion index for a new element in
/// a sorted collection that matches the rule defined by the predicate.
///
/// This variation is useful when comparing an element that
/// is of a different type than those already in the array.
/// - Parameter isOrderedAfter:
/// Return `true` if the new element should come after the one
/// provided in the closure, or `false` otherwise. For instance
/// `{ $0 < newElement }` to sort elements in increasing order.
/// - Returns: The index suitable for inserting the new element into
/// the array, or the first index of an existing element.
@inlinable
public func sortedInsertionIndex(
where isOrderedAfter: (Element) throws -> Bool
) rethrows -> Index {
var slice: SubSequence = self[...]
while !slice.isEmpty {
let middle = slice.index(slice.startIndex, offsetBy: slice.count/2)
if try isOrderedAfter(slice[middle]) {
slice = slice[index(after: middle)...]
} else {
slice = slice[..<middle]
}
}
return slice.startIndex
}
}
Since sometimes you don't care about the insertion index, but instead the first or last index that matches a given element, here are variations on the above that satisfy those requirements as well:
extension RandomAccessCollection {
@inlinable
public func sortedFirstIndex(
of element: Element
) -> Index? where Element: Comparable {
sortedFirstIndex(of: element, by: <)
}
@inlinable
public func sortedFirstIndex(
of element: Element,
by areInIncreasingOrder: (Element, Element) throws -> Bool
) rethrows -> Index? where Element: Comparable {
let insertionIndex = try sortedInsertionIndex(of: element, by: areInIncreasingOrder)
guard insertionIndex < endIndex, self[insertionIndex] == element else { return nil }
return insertionIndex
}
@inlinable
public func sortedLastIndex(
of element: Element
) -> Index? where Element: Comparable {
sortedLastIndex(of: element, by: <)
}
@inlinable
public func sortedLastIndex(
of element: Element,
by areInIncreasingOrder: (Element, Element) throws -> Bool
) rethrows -> Index? where Element: Comparable {
let insertionIndex = try sortedInsertionIndex(of: element) { try areInIncreasingOrder($1, $0) }
let finalIndex = index(insertionIndex, offsetBy: -1)
guard finalIndex >= startIndex, self[finalIndex] == element else { return nil }
return finalIndex
}
}