汎用的なバイナリサーチを実装して見た。
基本実装
まず本体
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
に準拠させることなくバイナリサーチが利用できます。