#　はじめに

以前、ソートされた配列の中に、とある要素が存在するかどうかを高速に検索する手法として、二分探索(ニブタン)を投稿した。
今回は、「一致する要素」以外にも、「最も近似する要素」等でもニブタンを利用して高速検索する例を投稿する。
なお、言語は、最近マイナーであることを身にしみて感じる「swift」で書いてます。AtCoderのABCコンテストのD問題以降だと、ACの解答例にswiftのコードが無いことがよくあるんだよね...悲しいね

おさらい

ニブタン（一致検索）のコードは次のとおり。なお、引数として渡す配列は、昇順でソート済みの前提。ニブタンの計算量はO(logN)だけど、ソートの計算量がO(NlogN)だから、ソートだけで全探索の計算量O(N)を超えてしまう。よって、ニブタンを繰り返し行うときに、ニブタンによる高速化は有効だけど、検索を一回だけ行うなら、ニブタンで無く全探索を行った方が全体の計算量は小さくなる。
以前作ったニブタンのコード(一致検索)をInt以外でも使えるように一般化し、また、配列の最小値未満または最大値超であれば、そこで探索をせずに終わるショートカットを追加した。

//二分探索(ニブタン)
func nib<T:Comparable>(_ vs:[T],_ num:T)->Bool{ // vsは昇順とする

    var l = 0 // 探索範囲の左端のindex
    var r = vs.count - 1 // 探索範囲の右端のindex

    if num < vs[l] || vs[r] < num {return false} // 配列の範囲外の時は終わり
    
    while(true) {
    
        let i = ( l + r ) / 2 // 探索範囲の中間地点
    
        if vs[i] == num {
            return true // 発見 
        } else if vs[i] > num { // インデックスiの要素の値が検索値numより大きいとき
            r = i - 1 //　探索範囲の右端をインデックスiの左隣とする
        } else { 
            l = i + 1  //　探索範囲の左端をインデックスiの右隣とする
        }
        
        if ( l > r ) {return false} // 発見できないまま、捜索範囲の左右が逆転
    }
}

let a = [1,3,5,7,9,11,15,20] // 昇順
print(nib(a,9)) // true
print(nib(a,10)) // false

let b = ["a","d","f","g","j","n"] // 昇順
print(nib(b,"d")) // true
print(nib(b,"e")) // false

近似値対応

次は、一致しない場合でも、近似値を返せるようにコードを書き換える。
上記の一致探索では、Comparableプロトコルに属する方を対象としたけど、近似値を求めるコードでは、引き算をする必要があるため、加算減算が可能なプロトコルAdditiveArithmeticを型制約に追加したうえで、以下の通りとした。

//ニブタン（近似値対応）
func nib<T:Comparable & AdditiveArithmetic>(_ vs:[T],_ num:T)->(Bool,T){ // vsは昇順とする

    var l = 0 // 探索範囲の左端のindex
    var r = vs.count - 1 // 探索範囲の右端のindex
    
    if num < vs[l] {return (false,vs[l])} // 配列の下限未満の時は、下限が近似値
    if vs[r] < num {return (false,vs[r])} // 配列の上限超の時は、上限が近似値
    
    while(true) {
    
        let i = ( l + r ) / 2 // 探索範囲の中間地点
    
        if vs[i] == num {
            return (true,num) // 発見 
        } else if vs[i] > num { // インデックスiの要素の値が検索値numより大きいとき
            r = i - 1 //　探索範囲の右端をインデックスiの左隣とする
        } else { 
            l = i + 1  //　探索範囲の左端をインデックスiの右隣とする
        }

        // 発見できないまま、捜索範囲の左右が逆転
        if ( l > r ) {
            if vs[l] - num >= num - vs[r] {return (false,vs[r])} else {return (false,vs[l])}   
        } 
    }
}

コードは上記のとおり、戻り値はタプル（一致する値の有無,近似値）とした。タプルの第1要素は元のニブタンと同じ、第2要素は近似値(第1要素がtrueの時は検索値)となる。
なお、近似値が2つあるとき(例えば、検索値が8で配列が[7,9]のとき、7と9は2つとも近似値)、小さい方が近似値となるコードになっている。
下記は実行例。

　// 入力
let a = [1,3,5,7,9,12,15,20] // 昇順
for i in 7...15 {
    print(i,nib(a,i))
}

// 出力
7 (true, 7)
8 (false, 7)
9 (true, 9)
10 (false, 9)
11 (false, 12)
12 (true, 12)
13 (false, 12)
14 (false, 15)
15 (true, 15)


// 入力
import Foundation
let b:[Decimal] = [-0.49,-0.35,-0.3,-0.03,0.1,0.27,0.4,0.51]
for i in stride(from: Decimal(-0.5), through: Decimal(0.5), by: Decimal(0.1)) {
    print(i,nib(b,i))
}

// 出力
-0.5 (false, -0.49)
-0.4 (false, -0.35)
-0.3 (true, -0.3)
-0.2 (false, -0.3)
-0.1 (false, -0.03)
0 (false, -0.03)
0.1 (true, 0.1)
0.2 (false, 0.27)
0.3 (false, 0.27)
0.4 (true, 0.4)
0.5 (false, 0.51)

まあ、お試しで作ってみたけど、使い道はあまりなさそう。次に紹介する、lower_boundとかの方が使える。というか、AtCoderでc++を使っている人は使いまくり。

c++の「lower_bound」,「upper_bound」について

競技プログラミングでc++を使っている人が多用している、ニブタン利用の検索機能にlower_boundとupper_boundがある。これは、ともに、検索結果の「位置(index)」を戻り値とする。有無だけを返す最初のニブタンと違い、位置情報なので、単なる有無情報(1bit情報)より価値が高いと思う。
なお、lowerとupperだから、真逆の機能かな？と思いきや、違いは「ほんのちょっと」なので、気を付ける必要がある。
例えば、配列の中身が、{10,10,30,30,30,40,70}の時に、検索値を「30」とした場合、
- lower_bound ：最初に検索値「以上」となるindexを返す。⇒ 「30以上」となる最初のindex「2」を返す。
- upper_bound ：最初に検索値「超」となるindexを返す。⇒ 「30超」となる最初のindex「5」を返す。
具体的にc++のコードで示すと

  auto a = {10,10,30,30,30,40,50,70,100,100};
  auto x = 30;

  auto lb = lower_bound(a.begin(), a.end(), x);
  auto ub = upper_bound(a.begin(), a.end(), x);
  
  cout << lb << " " << ub << endl; // 0x7fff732a7ca8 0x7fff732a7cb4
  cout << *lb << " " << *ub << endl; // 30 40
  
  cout << a.begin() << endl; // 0x7fff732a7ca0
  cout << ub - a.begin() << endl; // 5
  cout << ub - lb << endl; // 3

ちょ、待てよ！ lbが「2」じゃないじゃん！「0x7fff732a7ca8」って何だよ！
と思うのは、無理も無いですけど、スミマセン。わかりやすさを重視して、嘘つきました(ﾃﾍﾍﾟﾛ😛
でも、ある要素iとある要素jのindex同士の「差異」が計算可能なので、上記コードでも、ubのindexと先頭のindex(a.begin())の差により、もし、aが配列であった場合のubのindexとなる「5」(= ub - a.begin())を取得できるから、実質、lbは「2」、ubは「5」と言ってしまっても過言では無い、と言えなく無くも無い。
なお、検索値を100としてしまうと、upper_boundのindexが取得できない気がするけど、実際はindexは取得できる。でもインチキなindexだから、そのindexに格納されている値がおかしな数値となっている。じゃあ、なんでインチキindexを返すの？と思うけど、実は意味があるんです。後ほど説明します。

  auto x1 = 100;

  auto lb1 = lower_bound(a.begin(), a.end(), x1);
  auto ub1 = upper_bound(a.begin(), a.end(), x1);
  
  cout << lb1 << " " << ub1 << endl; // 0x7ffe1abe3f00 0x7ffe1abe3f08
  cout << *lb1 << " " << *ub1 << endl; // 100 4200016
  cout << ub1 - lb1 << endl; // 2

上記のとおり、lb1のindexは「0x7ffe1abe3f00」となり、そのindexにおける値*lb1は「100」となっている。これに対し、本来は条件を満たす要素が無いにも関わらず、indexは「0x7ffe1abe3f08」が得られている。そのindexにおける値*ub1については、「4200016」という意味の無い値となっているが、「ub1 - lb1」というindexの差が2であることにより、「検索値x1 = 100を満たす要素の数が2つある」という情報を得ることが出来た。よって、lower_bound、upper_boundの返り値(index)の条件として次が存在することが分かる
- 【追加条件】lower_bound,upper_boundの検索値「以上」または「超」となる値を含まないとき、最後の要素のindexの「次のindex」を返す。

swiftでlower_bound,upper_boundを実装する

データを保持する構造体として、配列を使います。だから、indexも0x7ffe1abe3f00みたいな解読不明な値（メモリアドレスなのかな？）を使わず、0スタートのindexを使えます。そして、配列内に検索値「以上」「超」の要素がないときのindexは[要素数]とします。（普通、配列の最後の要素の番号は[要素数-1]）
つくったコードは以下の通り。関数名は少し変えました。あと、whileループの条件を、r>lではなく、r-l>1に変えて、初期値もl=-1、r=vs.countにしています。これが、最もシンプルなコードっぽい。lower_bound、upper_boundを実装するときの定番の模様。

// lower_bound
func lbound<T:Comparable>(_ vs:[T],_ num:T)->Int{ // vsは昇順とする

    var l = -1 // 探索範囲の左端のindex
    var r = vs.count // 探索範囲の右端のindex
    
    while r - l > 1 {
        let m = (l + r) / 2
        if vs[m] >= num {r = m} else {l = m} //lower_bound
    }
    
    return r
}

// upper_bound
func ubound<T:Comparable>(_ vs:[T],_ num:T)->Int{ // vsは昇順とする

    var l = -1 // 探索範囲の左端のindex
    var r = vs.count // 探索範囲の右端のindex
    
    while r - l > 1 {
        let m = (l + r) / 2
        if vs[m] > num {r = m} else {l = m} //upper_bound
    }
    
    return r
}

動作チェックすると...うむ、いい感じ！

//入力
let a = [2,3,4] // 昇順

print("x i -- lower_bound")
for i in 0...5 {
    print(i,lbound(a,i))
}

print("x i -- upper_bound")
for i in 0...5 {
    print(i,ubound(a,i))
}

//出力
x i -- lower_bound
0 0
1 0
2 1
3 1
4 2
5 3

x i -- upper_bound
0 0
1 0
2 1
3 2
4 3
5 3

最初のc++のコード例も、同じ結果を再現できた!

let a = [10,10,30,30,30,40,50,70,100,100]
let x = 30

let lb = lbound(a, x)
let ub = ubound(a, x)
  
print(lb,ub) // 2 5
print(a[lb],a[ub]) // 30 40
print(ub - lb) // 3 -- 配列aの中の30の個数と一致

lboundとuboundを使って競技プログラミング問題を解いてみる

ニブタン（近似値対応）では、「検索値に1番近い数値」を探すことが出来たけど、「検索値とk番目に近い数値との距離」を見つけることが出来るだろうか？
そんな問題がABC364のd問題です。
ここで突然ですが、配列aに対して、検索値「x-d」「x+d」を使ってubound,lboundの値を求める関数「f(d)=ubound(a,x+d) - lbound(a,x-d)」はどんな意味を持つ数字だろうか？
ちなみに前項で例示したdがゼロの時、つまり、f(0)は配列中のxの個数でした。
f(d)は、「配列中の、x-d以上、x+d以下の要素の数」になります。例えば、a = [1,2,2,6,8]で、x=4のとき、f(0)=0,f(1)=0,f(2)=3,f(3)=4,f(4)=5となります。
ここで、当初の目的の「検索値とk番目に近い数値との距離」を考えると、f(d)>=kとなる最小のdを求めることと同じです。
おや、「□ >= ○ となる最小の△」ってどこかで見たことない？おぁ、lboundと同じじゃん？
△の位置にあるdが取り得る値は、0,1,2,3,4,....おお！indexっぽい！
□の位置にあるf(d)は単純増加関数になってるから、ソート済みの配列と同じじゃん！
でも、f(d)自体は配列ではないから、lbound関数を直接は使えないけど、lbound実装化の時のコードを使いまわせるね。
ということで、方向性が分かったところでコードを書いてみよう...の前に、どのような形で入力されるか確認する。

4 3 // 配列Aの要素の数4，クエリの数
-3 -1 5 6 // 配列Aの各要素（座標）の値
-2 3 // クエリ① 座標「-2」から3番目に近い配列Aの座標までの距離
2 1 // クエリ② 座標「2」から1番目に近い配列Aの座標までの距離
10 4 // クエリ③ 座標「10」から4番目に近い配列Aの座標までの距離

上記はプログラムを書かずとも分かる。
- クエリ①は、「−2」から「5」(3番目に近い座標)までの距離だから、出力は「7」となる。
- クエリ②は、「2」から「-1」or「5」(1番近い座標)までの距離だから、出力は「3」となる。
- クエリ③は、出力は「13」だね。
  -　それじゃ、コードを書いてみよう

let (N,Q) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]
var A = readLine()!.split(separator:" ").map{Int($0)!}
A.sort()

//クエリ
for _ in 0..<Q {
    let (b,k) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]

    //lboundのコードをパクり
    var l:Int = -1 // 探索範囲の左端のindex
    var r:Int = 200_000_000 // 探索範囲の右端のindex

    //lboundの対象として、昇順配列の代わりに単調関数
    func f(_ d:Int) -> Int {
        let ub = ubound(A,b + d)
        let lb = lbound(A,b - d)
        return ub - lb
    }

    //見た目は、配列を対象としたlower_boundの実装コードと同じ
    while r - l > 1 {
        let m = (l + r) / 2
        if f(m) >= k {r = m} else {l = m} //lower_bound
    }
    
    print(r)
}

提出したら、600ms以下でACだった！やったね！
とりあえず、swift教信者でも、lower_boundとupper_boundという武器を手に入れて、教プロで戦えることが分かったよ。

おわりに

最初からlower_boundやupper_boundが用意されているc++はズルいなぁ
Appleも、もう少し教プロ用のライブラリを充実させてほしいね。swift-collectionsはいつになったら、正式にswiftに取り込まれるんだろう。AtCoderでは使えるけど、paiza.ioとかで使えないから、不便だし、そもそもHeapはAtCoderでも使えなかったし。