Python3で実装問題解決力を鍛える！アルゴリズムとデータ構造【第６章】設計技法(4)：二分探索法

「問題解決力を鍛える！アルゴリズムとデータ構造」の章末問題をPython3で実装していくシリーズ。

大枠まとめ記事は、ここに。 blacktanktop.hatenablog.com

前回の記事は、【第５章】設計技法(3)：動的計画法

blacktanktop.hatenablog.com

今回は【第６章】設計技法(4)：二分探索法について扱う

問題解決力を鍛える!アルゴリズムとデータ構造 (KS情報科学専門書)

作者:大槻兼資
発売日: 2020/10/02
メディア: 単行本（ソフトカバー）

内容
【第６章】設計技法(4)：二分探索法の概要
章末問題
まとめ
終わりに

内容

Python3でやっている。
簡単に自分の理解を整理する用。
そこまで厳密ではない可能性が高い。
記事を書いたら、書評の方に記事を載せていく。
コード等はGitHubにおいた。 github.com

【第６章】設計技法(4)：二分探索法の概要

にぶたん。探索範囲を半減させながら探索を行うため、非常に高速に探索できる。

二分探索法のポイント

探索範囲を半分にさせていきながら解を求める手法。

一般化させていくと、応用範囲が多い。

章末問題

6.1 座標圧縮

入力した配列aが重複がない整数列として、その配列の要素が何番目に小さいかを出力する。これを座標圧縮という。

二分探索で使うbisectのポイント
1. bisect_left(a, x)とすれば $a_i \lt x$ になるような個数が返ってくる。また、その数を挿入した場合の挿入先のindexでもある。
2. bisect_right(a, x)とすれば $a_i \leq x$ になるような個数が返ってくる。また、その数を挿入した場合の挿入先のindexでもある。

今回は、bisect_leftを使って、二分探索の関数binary_search()を実装した。が、Pythonはbisect_leftという組み込み関数があるのでそれを使えばいい。

入力は以下を想定している。

12 43 7 15 9

from bisect import bisect_left

def main():
    def binary_search(list, n):
        # 左右のindexで考える
        l = 0
        h = len(list) - 1
        flag = False
        while l <= h:
            mid = (l+h) // 2
            tmp = list[mid]
            if n == tmp:
                flag = True
                break
            elif tmp < n:
                l = mid + 1
            elif tmp > n:
                h = mid - 1
        return mid
    
    a = [int(x) for x in input().split()]
    # 二分探索するためにソート
    a_s = sorted(a)
    print('ソートされたa', a_s)
    for i in a:
        print('i', i)
        print('（index的に）何番目に小さいか(自作関数)', binary_search(a_s, i))
        print('（index的に）何番目に小さいか（組み込み関数）', bisect_left(a_s, i))
if __name__ == '__main__':
    main()

6.2 Snuke Festival

出典は、この問題。 atcoder.jp

二分探索に気づいても、愚直に $a \lt b$ を探すために $a$ をforで回すと、その中で $b\lt c$ を探すために対象の $b$ をforで回すことになって、 $O(nlog(n))$ にはならない。

ではどうするかというと、 $b$ を中心に考える。
$b$ をforで回して、以下のようにやる。

$a \lt b$ を探す。
を探す。
- 上で「二分探索で使うbisectのポイント」を書いたが、bisectでは対象のリストの要素がある数値よりも小さい個数を返すか(bisect_left)、ある数値以下の個数を返す(bisect_right)
- なので、 $b \lt c$ を探すには、bisect_rightで $c \leq b$ を求める。
- その値を $n$ から引くことで求める。

入力例は

6
3 14 159 2 6 53
58 9 79 323 84 6
2643 383 2 79 50 288

b: 58
aの候補: [2, 3, 6, 14, 53]
cの候補: [79, 288, 383, 2643]
aの候補数 x cの候補数 20
b: 9
aの候補: [2, 3, 6]
cの候補: [50, 79, 288, 383, 2643]
aの候補数 x cの候補数 15
b: 79
aの候補: [2, 3, 6, 14, 53]
cの候補: [288, 383, 2643]
aの候補数 x cの候補数 15
b: 323
aの候補: [2, 3, 6, 14, 53, 159]
cの候補: [383, 2643]
aの候補数 x cの候補数 12
b: 84
aの候補: [2, 3, 6, 14, 53]
cの候補: [288, 383, 2643]
aの候補数 x cの候補数 15
b: 6
aの候補: [2, 3]
cの候補: [50, 79, 288, 383, 2643]
aの候補数 x cの候補数 10
なので、87が答え

from bisect import bisect_left, bisect_right

def main():
    n = int(input())
    a = [int(x) for x in input().split()]
    b = [int(x) for x in input().split()]
    c = [int(x) for x in input().split()]
    a_s = sorted(a)
    c_s = sorted(c)
    ans = 0
    for j in b:
        # a[i] < bの個数 
        a_b = bisect_left(a_s, j)
        print('b:', j)
        print('aの候補:', a_s[:a_b])
        # c[k] <= bの個数(求めたい方の逆)
        c_b = bisect_right(c_s, j)
        # b < c[k](全体の個数から引き算)
        b_c = n - c_b
        print('cの候補:', c_s[c_b:])
        print('aの候補数 x cの候補数', a_b * b_c)
        ans += a_b * b_c
    print('合計', ans)
if __name__ == '__main__':
    main()

6.3 ダーツ

問題の出典は第７回日本情報オリンピック本選の第３問

概要は、

$n$ 個の正の整数列 $a$ が与えられる。
重複を許して $a$ から4つ選んだ時の総和が $m$ を超えないようにする。
その時の総和の最大値は？
4つ選ぶ時に0もOK（何も選ばないもOK）
ただし計算量が $O(N^2log(N))$ で。

$O(N^2log(N))$ は結構厳しい。

愚直に、全てforで回すと $O(N^4))$
工夫して、6.2みたいに一つを固定する感じでやると $O(N^2log(N)))$
ではどうするか。
$a$ から2つ選んだ配列を作って、そこから1つずつ選ぶと考える。
まず、リストと値を与えたら、値を超えない最大値を得る関数をつくる。
bisectのポイントにも書いたが、bisect_right(a, x)とすれば $a_i \leq x$ になるような、その数を挿入した場合ので挿入先のindexを返してくるので値を超えない最大値のindexはそのindex-1となる。

from bisect import bisect_right
def listmax(l, x):
    return l[bisect_right(l, x)-1]

そして、あとは $a$ から2つ選んだ配列を作ってソート（配列 $q$ とする）
$q$ から1つずつ取り出して（値を $x$ とする）いって、 $x$ が $m$ を超えないなら、 $x + listmax(q, x)$ として、 $x$ ごとに最適な相手を見つけ、配列に加える。
最後にmaxをとる。入力は以下のような形を想定

from bisect import bisect_right

def main():
    # リストからx以下の最大値を得る
    def listmax(l, x):
        return l[bisect_right(l, x)-1]

    n, m = map(int, input().split())
    a = [0] + [int(input()) for _ in range(n)]
    q = sorted({i+j for i in a for j in a})
    print('aから重複（0も）許して２つ選んで足したリスト（qとする）', q)
    # m以下で抑えたいからx <= mの時に
    # qからm-x以下の最大値を得る
    # q = m-xのときはmぴったりになる
    qq = {x + listmax(q, m-x) for x in q if x <= m}
    print('qからmを超えないように2つ選んで足したリスト(qiごとに最大)', qq)
    ans = max(qq)
    print('最大値', ans)
if __name__ == '__main__':
    main()

6.4 Aggressive cows

保留

6.5 億マス計算

出典はこれ。

atcoder.jp

詳細をリンク先をみるとして、制約が結構きつい。

結果として $x$ 以下の数が $k$ 個以上あるような最小のxを求めるということである。　　
なので、まず $x$ 以下の積の数が $k$ 個以上あるかを判定する関数(check())を作る。
この問題では $a \times b \leq k$ は言い換えると $b \leq k//a$ (//は割った値の整数部分)となる。
あとは二分探索であり得る範囲の値、 $0 \leq x \leq 10^18+1$ で $x$ を探す。

入力は以下のような形を想定

from bisect import bisect_right

def main():
    n, k = map(int, input().split())
    A = [int(x) for x in input().split()]
    B = sorted([int(x) for x in input().split()])
    def check(x):
        # x以下の積の数がK個以上
        # a*b <= k　は b <= k//a
        cnt = 0
        for a in A:
            cnt += bisect_right(B, x // a)
        return cnt >= k
    # a. b の最大が10**9
    l = 0
    r = 10**18 + 1
    while r - l > 1:
        m = (r + l) // 2
        # m以下の積の数がK個以上ならmを右端に（現状のmより左側に答えが）
        if check(m):
            r = m
        # m以下の積の数がK個未満ならmを左端に（現状のmより右側に答えが）
        else:
            l = m
    # 全パターンを順番に並べた
    print('全パターンを順番に並べると')
    print(sorted(a * b for a in A for b in B))
    # 最後は右を答える
    print('小さい方から', k, '番目の値', r)
            
if __name__ == '__main__':
    main()

まとめ

二分探索はイメージしやすい。
複数の配列があるなど、探索する量が多い時は、複数の配列の１つを固定する、2つにまとめるなど、うまく工夫が必要。
問題に合わせてbisectで作ることができるチェック用の関数を作ると楽、例えば
- リストからx以下の最大値を得る関数
- $x$ 以下の積の数が $k$ 個以上あるかを判定する関数

終わりに