How do I insert an element at the correct position into a sorted array in Swift?

Question

NSArray has - (NSUInteger)indexOfObject:(id)obj inSortedRange:(NSRange)r options:(NSBinarySearchingOptions)opts usingComparator:(NSComparator)cmp to determine the insert position of a new object in a sorted array.

What is the best and high-performance way to do this in pure Swift?

Something along the lines of:

var myArray = ["b", "e", "d", "a"]
myArray.sort { $0 < $1 }

// myArray is now [a, b, d, e]

myArray.append("c")
myArray.sort { $0 < $1 }

// myArray is now [a, b, c, d, e]

Instead of appending the new element and then sorting the array, I would like to figure out the correct position and insert the element:

let index = [... how to calculate this index ??? ...]
myArray.insert("c", atIndex: index)

Solution 1

Here is a possible implementation in Swift using binary search (from http://rosettacode.org/wiki/Binary_search#Swift with slight modifications):

extension Array {
    func insertionIndexOf(_ elem: Element, isOrderedBefore: (Element, Element) -> Bool) -> Int {
        var lo = 0
        var hi = self.count - 1
        while lo <= hi {
            let mid = (lo + hi)/2
            if isOrderedBefore(self[mid], elem) {
                lo = mid + 1
            } else if isOrderedBefore(elem, self[mid]) {
                hi = mid - 1
            } else {
                return mid // found at position mid
            }
        }
        return lo // not found, would be inserted at position lo
    }
}

As with indexOfObject:inSortedRange:options:usingComparator: it is assumed that the array is sorted with respect to the comparator. It returns either (any) index of the element if the element is already present in the array, or the index where it can be inserted while preserving the order. This corresponds to the NSBinarySearchingInsertionIndex of the NSArray method.

Usage:

let newElement = "c"
let index = myArray.insertionIndexOf(newElement) { $0 < $1 } // Or: myArray.indexOf(c, <)
myArray.insert(newElement, at: index)

Solution 2

In swift 3 you can use index(where:):

var myArray = ["a", "b", "d", "e"]
let newElement = "c"
if let index = myArray.index(where: { $0 > newElement }) {
    myArray.insert(newElement, at: index)
}

Note that in this case you need to reverse the condition inside the closure (i.e. > instead of < for increasing elements in array), because the index you are interested in is the first element that does NOT match the predicate. Also, this method will return nil if the newly inserted element is going to be the last in the array (newElement = "z" in the example above.

For convenience, this can be wrapped to a separate function that will handle all these issues:

extension Collection {
    func insertionIndex(of element: Self.Iterator.Element,
                        using areInIncreasingOrder: (Self.Iterator.Element, Self.Iterator.Element) -> Bool) -> Index {
        return index(where: { !areInIncreasingOrder($0, element) }) ?? endIndex
    }
}

Usage:

var myArray = ["a", "b", "d", "e"]
let newElement = "c"
let index = myArray.insertionIndex(of: newElement, using: <)
myArray.insert(newElement, at: index)

Solution 3

According to WWDC 2018 Session 406: Swift Generics (Expanded) the binary search can be performed in a more efficient and even more generic way by slicing the collection object.

There are two considerable benefits of Slice:

1. A slice is always a subset of the original object without allocating additional memory.
2. All indices of the slice refer to the original object.
   For example if you slice an array of 5 objects let slice = array[2..<4] then slice.startIndex is 2 not 0.

RandomAccessCollection is a protocol (inherited from BidirectionalCollection) which a variety of structs / classes conform to

extension RandomAccessCollection where Element : Comparable {
    func insertionIndex(of value: Element) -> Index {
        var slice : SubSequence = self[...]

        while !slice.isEmpty {
            let middle = slice.index(slice.startIndex, offsetBy: slice.count / 2)
            if value < slice[middle] {
                slice = slice[..<middle]
            } else {
                slice = slice[index(after: middle)...]
            }
        }
        return slice.startIndex
    }
}

Example:

let array = [1, 2, 4, 7, 8]
let index = array.insertionIndex(of: 6) // 3

You can extend the function to check against a predicate closure instead of a single value

extension RandomAccessCollection { // the predicate version is not required to conform to Comparable
    func insertionIndex(for predicate: (Element) -> Bool) -> Index {
        var slice : SubSequence = self[...]

        while !slice.isEmpty {
            let middle = slice.index(slice.startIndex, offsetBy: slice.count / 2)
            if predicate(slice[middle]) {
                slice = slice[index(after: middle)...]
            } else {
                slice = slice[..<middle]
            }
        }
        return slice.startIndex
    }
}

Example:

struct Person { let name : String }

let array = [Person(name: "Adam"), Person(name: "Cynthia"), Person(name: "Nancy"), Person(name: "Tom")]
let index = array.insertionIndex{ $0.name < "Bruce" } // 1

Solution 4

If you know your array is sorted, you can use this method -- it will work with arrays of any type. It will traverse the whole array each time, so don't use this with large arrays - go for another data type if you have larger needs!

func insertSorted<T: Comparable>(inout seq: [T], newItem item: T) {
    let index = seq.reduce(0) { $1 < item ? $0 + 1 : $0 }
    seq.insert(item, atIndex: index)
}

var arr = [2, 4, 6, 8]
insertSorted(&arr, newItem: 5)
insertSorted(&arr, newItem: 3)
insertSorted(&arr, newItem: -3)
insertSorted(&arr, newItem: 11)
// [-3, 2, 3, 4, 5, 6, 8, 11]

Solution 5

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
    }
}

View more solutions