【競プロ】ソートされた2次元行列から要素を探し出す3通りの方法

本記事では，ソートされた行列から特定の要素を探し出すアルゴリズムを3通り紹介します．ソートされた行列とは

• ある要素より右側にある数字はその要素以上
• ある要素より下側にある数字はその要素以上

を満たすもので，例えば下の図のような行列です．

行列のサイズは$h*w$とします．この中に例えば「8」があるかどうかを判定し，あればそのインデックス（座標）を返すことが目標です（複数ある場合はどのインデックスを返してもいいものとします）．

まず1次元配列のおさらい

まず，ソートされた1次元の配列から特定の要素を探索する場合,全要素を一つずつ調べていると$O(n)$時間（$n$は配列の長さ）かかりますが，2分探索によって$O(\log n)$にできることが一般的に知られています．Pythonでは$bisect.bisect\_left()$などを使えば何も知らなくても勝手にできてしまいます．

import bisect

A = [1,4,5,5,7,8,10]

idx1 = bisect.bisect_left(A, 5)  # 5以上の最小のインデックスを返す
print(idx1)  # 2

idx2 = bisect.bisect(A, 5)  # 5より大きい最小のインデックスを返す，　bisect_rightも同様．
print(idx2)  # 4

ただこれでは元も子もないので，最も基本的な二分探索のやり方をおさらいします．配列を$A$，探したい値を$key$とします．まず配列の両端に二つのポインター$left = 0$と$right = n-1$を置きます．そして，$mid = \frac{left + right}{2}$ として

• $A[mid] < key$ ならば $left = mid+1$とする
• $A[mid] = key$ ならば発見できたのでそこで終了．
• $A[mid] > key$ ならば $right = mid-1$とする

を，$left <= right$である限り繰り返します．$left > right$になっても発見できなければ，$key$は$A$の中になかったということになります（判定条件や境界条件には色々と流儀がありますが，単純な二分探索はこれで十分動くはずです）．

def binary_search(A, key):
    n = len(A)
    left, right = 0, n - 1
    while left <= right:
        mid = (left+right) // 2
        if A[mid] < key:
            left = mid + 1
        elif A[mid] == key:
            return mid
        else:
            right = mid - 1

    return -1

二分探索の肝は，$mid$という一つの値を調べることで，その左右どちらかの領域にある要素は全て調べる必要がなくなるということです．これによって調べる要素の個数が減り，$O(\log n)$で探索することができるようになりました．これを生かして，行列に対してもなるべく探索する要素を減らせる方法を考えていきます．

2次元配列の探索1: O(h log w) の方法

まずはじめに思いつくのは，各行（または各列）に対して配列の２分探索を行う方法です．これは全ての行または列に対してfor文を回すだけで実行できます．

def search_matrix(self, A: List[List[int]], key: int) -> tuple:
    h = len(A)
    w = len(A[0])
        
    for i in range(h):  # 各行について配列の二分探索を行う
        left, right = 0, w - 1
        while left <= right:
            mid = (left+right) // 2
            if A[i][mid] < key:
                left = mid + 1
            elif A[i][mid] == key:
                return (i, mid)
            else:
                right = mid - 1
                    
        return (-1, -1)  # 発見できなかった場合
        

この計算量は$O(min(h \log w, w \log h))$となります．$h, w$の大小に合わせて，列を探索するか行を探索するかを決めれば良いです．これでも十分な速さで動くアルゴリズムと言えます．

2次元配列の探索2: O(h + w) の方法

ただ，これよりも速く，スマートな方法があります．例えば，現在行列の中の右上$(0, w-1)$にいるとします．もし，$key < A[0][w-1]$だった場合何が言えるでしょうか？行列中のある要素の下にある数は常にその要素以上であることが保証されています．既に$A[0][w-1]$が$key$より大きいのですから，これより下の数は必ず$key$より大きいことがわかり，もう探索する必要がありません．

よって次は$A[:][:w-2]$を探索します．この行列の右上$A[0][w-2]$に関して，今度は$key > A[0][w-2]$が成り立っていたとします．$A[0][w-2]$と同じ行でより左側にある要素は常に$A[0][w-2]$以下なので，これらに$key$は含まれません．よって次は$A[1:][:w-2]$の中を探せば十分でしょう．

これを繰り返していくと，いずれ$key$が見つかるか，もしくは探す行列がなくなって終わります．ここでは毎回探したい行列の右上の成分を$key$と比較することを繰り返していますが，これは次のような操作をしていると言い換えることができます．

• 行列の右上からスタートする
• 現在の要素$(i, j)$と$key$を比較
  • $A[i][j] < key$ならば次は$(i+1, j)$を探索
  • $A[i][j] == key$ならば発見，終了
  • $A[i][j] > key$ならば次は$(i, j-1)$を探索
• もし$i >= h$または$j < 0$になれば発見できず終了

これはつまり，行列の右上から左下の方へと斜めに要素を探索していくだけで全要素を探索できるということです．

def search_matrix(self, A: List[List[int]], key: int) -> tuple:
    h = len(A)
    w = len(A[0])
        
    i, j = 0, w - 1  # 右上からスタート
        
    while i < w and j >= 0:  # 行列の範囲から出てしまうまで続ける
        if matrix[row][column] < key:
            i += 1
        elif A[i][j] == key:
            return (i, j)
        else:
            j -= 1
                
    return (-1, -1)  # 発見できなかった場合

この場合，最も時間がかかるケースは右上から左下まで移動するケースなので，時間計算量は$O(h + w)$になります．

2次元配列の探索3: O(h log (w/h)) の方法

方法2でも十分速いのですが，なんというか二分探索っぽくありません．2次元でも二分探索する方法はないのでしょうか？

実はあります．
　
ある要素$A[i][j]$が$key$より小さかった場合に何が言えるかもう一度考えてみましょう．$A[i][j]$の左側の要素と上側の要素は$key$以下です．それだけでなく，それらに囲まれた左上の要素全てが$key$より小さくなるはずです．逆に$A[i][j] > key$だった場合，$A[i][j]$を右上とした矩形領域の要素は全て$key$より大きいことになります．

例えば行列の$i$行目で二分探索を行い，求める$key$が見つからなかったとします．二分探索が終わった後，2つのポインタは$right+1 = left$という位置関係にあり，$right$は$key$より小さいもののうち最大の要素，$left$は$key$より大きいもののうち最小の要素にあります．すると，$A[i][right]$を右下とした行列の要素は全て$key$より小さく，$A[i][left]$を左上とした行列の要素は全て$key$より大きいことがわかります．

このエリアに$key$はありません．よって残された右上と左下の領域のみを探索すれば良いことになります．この2つの領域は

• 右上領域: 縦 $0$ ~ $i-1$, 横 $left$ ~ $w-1$
• 左下領域: 縦 $i+1$ ~ $h-1$, 横 $0$ ~ $right$

となります．

それが求まれば，次は2つの領域それぞれについて同じように探索していきます．小行列の範囲がなくなれば探索終了です．探索される小行列の数は再帰的に増えていきますが，その都度探索しなくていい領域も分かってくるので，全体として探索する領域は少なくなり，効率的に計算ができます．なお二分探索は領域のちょうど真ん中の行で行えば十分だと思われます．

def search_matrix(self, A: List[List[int]], key: int) -> tuple:
    h = len(A)
    w = len(A[0])

    def search_subarea(top_row, left_col, bottom_row, right_col):  # 部分行列を探索する関数
        if top_row > bottom_row or left_col > right_col:  # 領域が存在しなければ探索しない
            return (-1, -1)

        mid_row = (top_row + bottom_row) // 2  # 真ん中の行
        left, right = left_col, right_col
        while left <= right:  # 二分探索
            mid_col = (left+right) // 2
            if A[mid_row][mid_col] < key:
                left = mid_col + 1
            elif A[mid_row][mid_col] == key:
                return (mid_row, mid_col)
            else:
                right = mid_col - 1

        # 発見できなければ2つの部分行列を探索
        upper_right = search_subarea(top_row, left, mid_row - 1, right_col)
        lower_left = search_subarea(mid_row + 1, left_col, bottom_row, right)
        
        if upper_right != (-1, -1):
            return upper_right
        elif lower_left != (-1, -1):
            return lower_left
        else:
            return (-1, -1)

    return search_subarea(0, 0, h-1, w-1)
    

上のコードでは、部分行列の中を再帰的に探索する$search\_subarea$関数を動かしています．入力の$(top\_row, left\_col, bottom\_row, right\_col)$はそれぞれ，探索する行列の上端，左端，下端，右端の座標を表しています．中央の行である$mid\_row$に対して二分探索を行い，$key$を発見できなければ$left, right$を元に2つの小行列を再帰的に探索していきます．

この方法の計算量を考えるのは少々複雑です．縦$h$，横$w$（$h < w$とする）の行列に対しては，まず$log(w)$で$mid\_col$を探索し，その後最悪ケースでちょうど半分の地点で分割された2つの小行列を探索するため，計算時間は漸化的に　$T(h, w) = \log w + 2 * T(h/2, w/2)$　と書けます．stack-overflowの記事によれば，この漸化式を解くとどうやら　$O(h * \log \frac{w}{h})$　のオーダーになるらしいです．（私もよくわかっていないので教えてください...）

結局どれが速いのか

• 方法1: $O(h \log w)$　（$h < w$のとき）
• 方法2: $O(h + w)$
• 方法3: $O(h * \log \frac{w}{h})$　（$h < w$のとき）

結局どれが速いのでしょうか？まず方法1と方法3に関しては，時間計算量の式を見るに方法3の方が常に速いと言えます．方法3はうまく探索範囲を絞ることで二分探索する行の範囲を減らしているので，直感的にも方法3の方が速いのは納得がいきます．

問題は方法2と方法3の比較ですが，これは場合によります．具体的には，$h$と$w$の比によってどちらが適しているかが変化します．方法1の時間計算量は長辺の長さがボトルネックになる一方，方法3の時間計算量はlogの外に出ている短辺の長さがボトルネックになります．厳密な議論ができず申し訳ないのですが，stack-overflowの記事によれば，$100 \times 100$の行列では方法1の方が速く，$10
\times 1000$の行列では方法3の方が速くなるようです．

例えば極端な例として$1 \times 10000$の行列を考えた場合，方法2は端から全要素を探索するような挙動になってしまうため，二分探索を行う方法3のほうが明らかに速いです．

実装面を考えると方法2がシンプルで良いですが，縦横比が極端に違うような配列では方法3を使う方が良いのだと思われます．

まとめ

2次元のソートされた行列から要素を探す3通りの方法を紹介しました．それぞれ長所と短所があることがわかったので，今度街でソートされた行列を見かけた際には，状況に応じてアルゴリズムを使い分けられると周りの友達に一目置かれること間違いなしです．