Swift: Binary search for standard array?

Here's a generic way to use binary search:

func binarySearch<T:Comparable>(_ inputArr:Array<T>, _ searchItem: T) -> Int? {
    var lowerIndex = 0
    var upperIndex = inputArr.count - 1

    while (true) {
        let currentIndex = (lowerIndex + upperIndex)/2
        if(inputArr[currentIndex] == searchItem) {
            return currentIndex
        } else if (lowerIndex > upperIndex) {
            return nil
        } else {
            if (inputArr[currentIndex] > searchItem) {
                upperIndex = currentIndex - 1
            } else {
                lowerIndex = currentIndex + 1
            }
        }
    }
}

var myArray = [1,2,3,4,5,6,7,9,10]
if let searchIndex = binarySearch(myArray, 5) {
    print("Element found on index: \(searchIndex)")
}

Here's my favorite implementation of binary search. It's useful not only for finding the element but also for finding the insertion index. Details about assumed sorting order (ascending or descending) and behavior with respect to equal elements are controlled by providing a corresponding predicate (e.g. { $0 < x } vs { $0 > x } vs { $0 <= x } vs { $0 >= x }). The comment unambiguously says what exactly does it do.

extension RandomAccessCollection {
    /// Finds such index N that predicate is true for all elements up to
    /// but not including the index N, and is false for all elements
    /// starting with index N.
    /// Behavior is undefined if there is no such N.
    func binarySearch(predicate: (Element) -> Bool) -> Index {
        var low = startIndex
        var high = endIndex
        while low != high {
            let mid = index(low, offsetBy: distance(from: low, to: high)/2)
            if predicate(self[mid]) {
                low = index(after: mid)
            } else {
                high = mid
            }
        }
        return low
    }
}

Example usage:

(0 ..< 778).binarySearch { $0 < 145 } // 145

I use an extension on Indexable implementing indexOfFirstObjectPassingTest.

  • It takes a test predicate, and returns the index of the first element to pass the test.
  • If there is no such index, then it returns endIndex of the Indexable.
  • If the Indexable is empty, you get the endIndex.

Example

let a = [1,2,3,4]

a.map{$0>=3}
// returns [false, false, true, true]

a.indexOfFirstObjectPassingTest {$0>=3}
// returns 2

Important

You need to ensure test never returns in false for any index after an index it has said true for. This is equivalent to the usual precondition that binary search requires your data to be in order.

Specifically, you must not do a.indexOfFirstObjectPassingTest {$0==3}. This will not work correctly.

Why?

indexOfFirstObjectPassingTest is useful because it lets you find ranges of stuff in your data. By adjusting the test, you can find the lower and upper limits of "stuff".

Here's some data:

let a = [1,1,1, 2,2,2,2, 3, 4, 5]

We can find the Range of all the 2s like this…

let firstOf2s = a.indexOfFirstObjectPassingTest({$0>=2})
let endOf2s = a.indexOfFirstObjectPassingTest({$0>2})
let rangeOf2s = firstOf2s..<endOf2s
  • If there are no 2s in the data, we'll get back an empty range, and we don't need any special handling.
  • Provided there are 2s, we'll find all of them.

As an example, I use this in an implementation of layoutAttributesForElementsInRect. My UICollectionViewCells are stored sorted vertically in an array. It's easy to write a pair of calls that will find all cells that are within a particular rectangle and exclude any others.

Code

extension Indexable {
  func indexOfFirstObjectPassingTest( test: (Self._Element -> Bool) ) -> Self.Index {
    var searchRange = startIndex..<endIndex

    while searchRange.count > 0 {
      let testIndex: Index = searchRange.startIndex.advancedBy((searchRange.count-1) / 2)
      let passesTest: Bool = test(self[testIndex])

      if(searchRange.count == 1) {
        return passesTest ? searchRange.startIndex : endIndex
      }

      if(passesTest) {
        searchRange.endIndex = testIndex.advancedBy(1)
      }
      else {
        searchRange.startIndex = testIndex.advancedBy(1)
      }
    }

    return endIndex
  }
}

Disclaimer & Caution

I have about 6 years of iOS experience, 10 of Objective C, and >18 programming generally…

…But I'm on day 3 of Swift :-)

  1. I've used an extension on the Indexable protocol. This might be stupid approach – feedback welcomed.
  2. Binary searches are notoriously hard to correctly code. You really should read that link to find out just how common mistakes in their implementation are, but here is an extract:

When Jon Bentley assigned it as a problem in a course for professional programmers, he found that an astounding ninety percent failed to code a binary search correctly after several hours of working on it, and another study shows that accurate code for it is only found in five out of twenty textbooks. Furthermore, Bentley's own implementation of binary search, published in his 1986 book Programming Pearls, contains an error that remained undetected for over twenty years.

Given that last point, here are test for this code. They pass. They are unlikely to be exhaustive – so there may certainly still be errors. The tests are not guaranteed to actually be correct! There are no tests for the tests.

Tests

class BinarySearchTest: XCTestCase {

  func testCantFind() {
    XCTAssertEqual([].indexOfFirstObjectPassingTest {(_: Int) -> Bool in false}, 0)
    XCTAssertEqual([1].indexOfFirstObjectPassingTest {(_: Int) -> Bool in false}, 1)
    XCTAssertEqual([1,2].indexOfFirstObjectPassingTest {(_: Int) -> Bool in false}, 2)
    XCTAssertEqual([1,2,3].indexOfFirstObjectPassingTest {(_: Int) -> Bool in false}, 3)
    XCTAssertEqual([1,2,3,4].indexOfFirstObjectPassingTest {(_: Int) -> Bool in false}, 4)
  }

  func testAlwaysFirst() {
    XCTAssertEqual([].indexOfFirstObjectPassingTest {(_: Int) -> Bool in true}, 0)
    XCTAssertEqual([1].indexOfFirstObjectPassingTest {(_: Int) -> Bool in true}, 0)
    XCTAssertEqual([1,2].indexOfFirstObjectPassingTest {(_: Int) -> Bool in true}, 0)
    XCTAssertEqual([1,2,3].indexOfFirstObjectPassingTest {(_: Int) -> Bool in true}, 0)
    XCTAssertEqual([1,2,3,4].indexOfFirstObjectPassingTest {(_: Int) -> Bool in true}, 0)
  }

  func testFirstMatch() {
    XCTAssertEqual([1].indexOfFirstObjectPassingTest {1<=$0}, 0)
    XCTAssertEqual([0,1].indexOfFirstObjectPassingTest {1<=$0}, 1)
    XCTAssertEqual([1,2].indexOfFirstObjectPassingTest {1<=$0}, 0)
    XCTAssertEqual([0,1,2].indexOfFirstObjectPassingTest {1<=$0}, 1)
  }

  func testLots() {
    let a = Array(0..<1000)
    for i in a.indices {
      XCTAssertEqual(a.indexOfFirstObjectPassingTest({Int(i)<=$0}), i)
    }
  }
}


extension ArraySlice where Element: Comparable {
    func binarySearch(_ value: Element) -> Int? {
        guard !isEmpty else { return nil }

        let midIndex = (startIndex + endIndex) / 2
        if value == self[midIndex] {
            return midIndex
        } else if value > self[midIndex] {
            return self[(midIndex + 1)...].binarySearch(value)
        } else {
            return self[..<midIndex].binarySearch(value)
        }
    }
}

extension Array where Element: Comparable {
    func binarySearch(_ value: Element) -> Int? {
        return self[0...].binarySearch(value)
    }
}

This is, in my opinion, very readable and leverages the fact that Swift's ArraySlice is a view on Array and retains the same indexes as the original Array with which it shares the storage so, in absence of mutations (like in this case), it is therefore very efficient.

Tags:

Arrays

Types

Binary Search

Swift

Recent Posts