[Swift] ソート済み配列をバイナリサーチ

汎用的なバイナリサーチを実装して見た。

基本実装

まず本体

extension MutableCollection where IndexDistance == Int {
    private func bsearch(min: Int, max: Int, comparator: (Iterator.Element) -> ComparisonResult) -> Iterator.Element? {
        if max < min { return nil }
        let current = min + (max - min) / 2
        let v = self[self.index(self.startIndex, offsetBy: current)]
        let compRes = comparator(v)

        if compRes == .orderedSame { return v }

        let newMin = (compRes == .orderedAscending) ? current + 1 : min
        let newMax = (compRes == .orderedDescending) ? current - 1 : max
        return bsearch(min: newMin, max: newMax, comparator: comparator)
    }

    func binarySearch(comparator: (Iterator.Element) -> ComparisonResult) -> Iterator.Element? {
        return bsearch(min: 0, max: self.count - 1, comparator: comparator)
    }
}

Arrayでいいと思うんだけど、よくわからないのでsorted()を持ってるMutableCollectionを拡張しました。
間違えてる場合は指摘してくださるとありがたいです。
（いつか実装されるはずの末尾再帰の最適化のために末尾再帰にしています）

使用法は

let original = [8, 5, 9, 10, 66, 4, 3, 6, 12, 13, 16, 27, 58, 3]
let sortedArray = original.sorted()
let target = 5
sortedArray.binarySearch {
    if $0 == target { return .orderedSame }
    if $0 < target { return .orderedAscending }
    return .orderedDescending
} // 5

のようになります。

補助的実装

クロージャの中がゴチャゴチャしていてすごくわかりにくいです。
なのでComparisonResultを返す演算子を作っちゃいます。

infix operator ==? : ComparisonPrecedence
func ==? <T: Comparable> (lhs: T, rhs: T) -> ComparisonResult {
    if lhs == rhs { return .orderedSame }
    if lhs < rhs { return .orderedAscending }
    return .orderedDescending
}

ここでは==?を使いましたがお好みで。
比較対象がComparableに準拠していれば使えます。

これを使うと先のバイナリサーチはこうなります。

sortedArray.binarySearch { $0 ==? target }  // 5

一気に見やすくわかりやすくなりました。演算子偉大です。

補足

比較部分が外部に出ているので要素自体がComparableに準拠していなくても使えます。

struct Member {
    let id: Int
    let name: String
}

let members = sortedArray.map { Member(id: $0, name: "") }

members.binarySearch { $0.id ==? target }   // Member(id: 5, name: "")

比較対象がIntのため、構造体Member自体をComparableに準拠させることなくバイナリサーチが利用できます。