An array that keeps its elements sorted at all times.

SortedArray.swift

/// An array that keeps its elements sorted at all times.
public struct SortedArray<Element> {
    // Not sure if it's a good idea to use `ArraySlice` as the backing store. It lets me make SortedArray.SubSequence == SortedArray, but the price you pay for that is that one small slice, if stored permanently and not just locally inside a function, can easily retain a much larger collection, and this is hard to notice by the developer.
    fileprivate var _elements: ArraySlice<Element>

    public typealias Comparator<A> = (A, A) -> Bool

    /// The predicate that determines the array's sort order.
    fileprivate let areInIncreasingOrder: Comparator<Element>

    /// Initializes an empty array.
    ///
    /// - Parameter comparator: The comparison predicate the array should use to sort its elements.
    public init(comparator: @escaping Comparator<Element>) {
        _elements = []
        areInIncreasingOrder = comparator
    }

    /// Initializes the array with a sequence of unsorted elements and a comparison predicate.
    public init<S: Sequence>(unsorted: S, comparator: @escaping Comparator<Element>) where S.Iterator.Element == Element {
        let sorted = unsorted.sorted(by: comparator)
        _elements = sorted[sorted.startIndex ..< sorted.endIndex]
        areInIncreasingOrder = comparator
    }

    /// Initializes the array with a sequence that is already sorted according to the given comparison predicate.
    ///
    /// This is faster than `init(unsorted:comparator:)` because the elements don't have to sorted again.
    ///
    /// - Precondition: `sorted` is sorted according to the given comparison predicate. If you violate this condition, the behavior is undefined.
    public init<S: Sequence>(sorted: S, comparator: @escaping Comparator<Element>) where S.Iterator.Element == Element {
        _elements = ArraySlice(sorted)
        areInIncreasingOrder = comparator
    }

    /// Inserts a new element into the array, preserving the sort order.
    ///
    /// - Returns: the index where the new element was inserted.
    /// - Complexity: O(_n_) where _n_ is the size of the array. O(_log n_) if the new element can be appended, i.e. if it is ordered last in the resulting array.
    @discardableResult
    public mutating func insert(_ newElement: Element) -> Index {
        let index = insertionIndex(for: newElement)
        // Performance optimization: we can simply append in O(1) if the element is to be inserted at the end. insert(_:at:) probably already performs this optimization internally, but I'm not 100% sure.
        switch index {
        case endIndex: _elements.append(newElement)
        default: _elements.insert(newElement, at: index)
        }
        return index
    }

    /// Inserts all elements from `elements` into `self`, preserving the sort order.
    ///
    /// This can be faster than inserting the individual elements one after another because we only need to re-sort once.
    ///
    /// - Complexity: O(_n * log(n)_) where _n_ is the size of the resulting array.
    public mutating func insert<S: Sequence>(contentsOf newElements: S) where S.Iterator.Element == Element {
        _elements.append(contentsOf: newElements)
        _elements.sort(by: areInIncreasingOrder)
    }
}

fileprivate extension SortedArray {
    init(sortedSlice sorted: ArraySlice<Element>, comparator: @escaping Comparator<Element>) {
        _elements = sorted
        areInIncreasingOrder = comparator
    }
}

extension SortedArray where Element: Comparable {
    /// Initializes an empty sorted array. Uses `<` as the comparison predicate.
    public init() {
        self.init(comparator: <)
    }

    /// Initializes the array with a sequence of unsorted elements. Uses `<` as the comparison predicate.
    public init<S: Sequence>(unsorted: S) where S.Iterator.Element == Element {
        self.init(unsorted: unsorted, comparator: <)
    }

    /// Initializes the array with a sequence that is already sorted according to the `<` comparison predicate. Uses `<` as the comparison predicate.
    ///
    /// This is faster than `init(unsorted:)` because the elements don't have to sorted again.
    ///
    /// - Precondition: `sorted` is sorted according to the `<` predicate. If you violate this condition, the behavior is undefined.
    public init<S: Sequence>(sorted: S) where S.Iterator.Element == Element {
        self.init(sorted: sorted, comparator: <)
    }
}

extension SortedArray: RandomAccessCollection {
    public typealias Index = Int
    public typealias Indices = DefaultRandomAccessIndices<SortedArray<Element>>
    public typealias SubSequence = SortedArray<Element>

    public var startIndex: Index { return _elements.startIndex }
    public var endIndex: Index { return _elements.endIndex }

    public func index(after i: Index) -> Index {
        return _elements.index(after: i)
    }

    public func index(before i: Index) -> Index {
        return _elements.index(before: i)
    }

    public func index(_ i: Index, offsetBy n: Index) -> Index {
        return _elements.index(i, offsetBy: n)
    }

    public func distance(from start: Index, to end: Index) -> Index {
        return _elements.distance(from: start, to: end)
    }

    public subscript(position: Index) -> Element {
        return _elements[position]
    }

    public subscript(range: Range<Index>) -> SubSequence {
        return SortedArray(sortedSlice: _elements[range], comparator: areInIncreasingOrder)
    }
}


// MARK: - More efficient overrides of default implementations or implementations that need fewer constraints than the default implementations.
extension SortedArray {
    /// Returns the first index where the specified value appears in the collection.
    ///
    /// - Complexity: O(_log(n)_), where _n_ is the size of the array.
    public func index(of element: Element) -> Index? {
        switch search(for: element) {
        case let .found(at: index): return index
        case .notFound(insertAt: _): return nil
        }
    }

    /// Returns a Boolean value indicating whether the sequence contains the given element.
    ///
    /// - Complexity: O(_log(n)_), where _n_ is the size of the array.
    public func contains(_ element: Element) -> Bool {
        return index(of: element) != nil
    }

    /// Returns the minimum element in the sequence.
    ///
    /// - Complexity: O(1).
    @warn_unqualified_access
    public func min() -> Element? {
        return _elements.first
    }

    /// Returns the maximum element in the sequence.
    ///
    /// - Complexity: O(1).
    @warn_unqualified_access
    public func max() -> Element? {
        return _elements.last
    }
}

// MARK: - Binary search
extension SortedArray {
    /// The index where `newElement` should be inserted to preserve the array's sort order.
    fileprivate func insertionIndex(for newElement: Element) -> Index {
        switch search(for: newElement) {
        case let .found(at: index): return index
        case let .notFound(insertAt: index): return index
        }
    }
}

fileprivate enum Match<Index: Comparable> {
    case found(at: Index)
    case notFound(insertAt: Index)
}

extension SortedArray {
    /// Searches the array for `newElement` using binary search.
    ///
    /// - Returns: If `newElement` is in the array, returns `.found(at: index)` where `index` is the index of the element in the array. If `newElement` is not in the array, returns `.notFound(insertAt: index)` where `index` is the index where the element should be inserted to preserve the sort order. If the array contains multiple elements that are equal to `newElement`, there is no guarantee which of these is found.
    /// - Complexity: O(_log(n)_), where _n_ is the size of the array.
    fileprivate func search(for newElement: Element) -> Match<Index> {
        guard !isEmpty else { return .notFound(insertAt: endIndex) }
        var left = startIndex
        var right = index(before: endIndex)

        while left <= right {
            let dist = distance(from: left, to: right)
            let mid = index(left, offsetBy: dist/2)
            let candidate = self[mid]

            if areInIncreasingOrder(candidate, newElement) {
                left = index(after: mid)
            } else if areInIncreasingOrder(newElement, candidate) {
                right = index(before: mid)
            } else {
                // If neither element comes before the other, they _must_ be
                // equal, per the strict ordering requirement of `areInIncreasingOrder`.
                return .found(at: mid)
            }
        }
        // Not found. left is the index where this element should be placed if it were inserted.
        return .notFound(insertAt: left)
    }
}

SortedArrayTests.swift

import Toolbox
import XCTest

func assertEqualElements<S1, S2>(_ expression1: @autoclosure () throws -> S1, _ expression2: @autoclosure () throws -> S2, _ message: @autoclosure () -> String = "", file: StaticString = #file, line: UInt = #line)
    where S1: Sequence, S2: Sequence, S1.Iterator.Element == S2.Iterator.Element, S1.Iterator.Element: Equatable {

    // This should give a better error message than using XCTAssert(try expression1().elementsEqual(expression2()), ...)
    XCTAssertEqual(Array(try expression1()), Array(try expression2()), message, file: file, line: line)
}

class SortedArrayTests: XCTestCase {
    override func setUp() {
        super.setUp()
    }
    
    override func tearDown() {
        super.tearDown()
    }
    
    func testInitUnsortedSorts() {
        let sut = SortedArray(unsorted: [3,4,2,1], comparator: <)
        assertEqualElements(sut, [1,2,3,4])
    }

    func testInitSortedDoesntResort() {
        // Warning: this is not a valid way to create a SortedArray
        let sut = SortedArray(sorted: [3,2,1])
        assertEqualElements(Array(sut), [3,2,1])
    }

    func testSortedArrayCanUseArbitraryComparisonPredicate() {
        struct Person {
            var firstName: String
            var lastName: String
        }
        let a = Person(firstName: "A", lastName: "Smith")
        let b = Person(firstName: "B", lastName: "Jones")
        let c = Person(firstName: "C", lastName: "Lewis")

        var sut = SortedArray<Person> { $0.firstName > $1.firstName }
        sut.insert(contentsOf: [b,a,c])
        assertEqualElements(sut.map { $0.firstName }, ["C","B","A"])
    }

    func testConvenienceInitsUseLessThan() {
        let sut = SortedArray(unsorted: ["a","c","b"])
        assertEqualElements(sut, ["a","b","c"])
    }

    func testInsertAtBeginningPreservesSortOrder() {
        var sut = SortedArray(unsorted: 1...3)
        sut.insert(0)
        assertEqualElements(sut, [0,1,2,3])
    }

    func testInsertInMiddlePreservesSortOrder() {
        var sut = SortedArray(unsorted: 1...5)
        sut.insert(4)
        assertEqualElements(sut, [1,2,3,4,4,5])
    }

    func testInsertAtEndPreservesSortOrder() {
        var sut = SortedArray(unsorted: 1...3)
        sut.insert(5)
        assertEqualElements(sut, [1,2,3,5])
    }

    func testInsertAtBeginningReturnsInsertionIndex() {
        var sut = SortedArray(unsorted: [1,2,3])
        let index = sut.insert(0)
        XCTAssertEqual(index, 0)
    }

    func testInsertInMiddleReturnsInsertionIndex() {
        var sut = SortedArray(unsorted: [1,2,3,5])
        let index = sut.insert(4)
        XCTAssertEqual(index, 3)
    }

    func testInsertAtEndReturnsInsertionIndex() {
        var sut = SortedArray(unsorted: [1,2,3])
        let index = sut.insert(100)
        XCTAssertEqual(index, 3)
    }

    func testInsertInEmptyArrayReturnsInsertionIndex() {
        var sut = SortedArray<Int>()
        let index = sut.insert(10)
        XCTAssertEqual(index, 0)
    }

    func testInsertEqualElementReturnsCorrectInsertionIndex() {
        var sut = SortedArray(unsorted: [3,1,0,2,1])
        let index = sut.insert(1)
        XCTAssert(index == 1 || index == 2 || index == 3)
    }

    func testInsertContentsOfPreservesSortOrder() {
        var sut = SortedArray(unsorted: [10,9,8])
        sut.insert(contentsOf: (7...11).reversed())
        assertEqualElements(sut, [7,8,8,9,9,10,10,11])
    }

    func testIndexOfFindsElementInMiddle() {
        let sut = SortedArray(unsorted: ["a","z","r","k"])
        let index = sut.index(of: "k")
        XCTAssertEqual(index, 1)
    }

    func testIndexOfFindsFirstElement() {
        let sut = SortedArray(sorted: 1..<10)
        let index = sut.index(of: 1)
        XCTAssertEqual(index, 0)
    }

    func testIndexOfFindsLastElement() {
        let sut = SortedArray(sorted: 1..<10)
        let index = sut.index(of: 9)
        XCTAssertEqual(index, 8)
    }

    func testIndexOfReturnsNilWhenNotFound() {
        let sut = SortedArray(unsorted: "Hello World".characters)
        let index = sut.index(of: "h")
        XCTAssertNil(index)
    }

    func testIndexOfReturnsFirstMatchForDuplicates() {
        let sut = SortedArray(unsorted: "abcabcabc".characters)
        let index = sut.index(of: "c")
        XCTAssertEqual(index, 6)
    }

    func testSubSequenceIsSortedArray() {
        let sut = SortedArray(sorted: 1...10)
        let slice = sut.dropFirst(5)
        XCTAssert(slice as Any is SortedArray<Int>)
        assertEqualElements(slice, 6...10)
    }

    func testsContains() {
        let sut = SortedArray(unsorted: "Lorem ipsum".characters)
        XCTAssertTrue(sut.contains(" "))
        XCTAssertFalse(sut.contains("a"))
    }

    func testMin() {
        let sut = SortedArray(unsorted: -10...10)
        XCTAssertEqual(sut.min(), -10)
    }

    func testMax() {
        let sut = SortedArray(unsorted: -10...(-1))
        XCTAssertEqual(sut.max(), -1)
    }
}