【AtCoder】Pythonで二分探索

2021.10.01

2024.03.24

競技プログラミング

AtCoderPython二分探索

はじめに

Pythonで二分探索を実装するためのテンプレとAtCoderの問題をまとめました。

Pythonで二分探索

二分探索とは、検索する問題を半分に分割しながら解を探しだすアルゴリズムです。

Pythonでソートされたリストに対して、ソートされた状態を保ちながら挿入・挿入できる場所を求める場合はbisectというライブラリが便利です。下記の記事で解説しています。

【AtCoder】bisectでリストを二分探索する

はじめに bisectの使い方とbisectを使うAtCoderの問題をまとめました。 bis

本記事で紹介する二分探索はbisectを利用せず、ある範囲から解を探すための実装になります。

二分探索のテンプレ

bisectを利用しない二分探索のテンプレは以下になります。

1def is_ok(mid):
2    """二分探索中の判定
3
4    Parameters
5    ----------
6    mid : int
7
8    Returns
9    -------
10    resutl : bool
11    """
12    if mid > 0:
13        return True
14    else:
15        return False
16
17
18def binary_search(ok, ng):
19    """二分探索
20
21    Parameters
22    ----------
23    ok : int
24    ng : int
25
26    Returns
27    -------
28    ok : int
29    """
30    while ng - ok > 1:
31        mid = (ok + ng) // 2
32        if is_ok(mid):
33            ok = mid
34        else:
35            ng = mid
36    return ok

binary_search関数で検索する範囲を半分にしながら解を探します。

is_ok関数の判定基準を修正することでいろんなパターンに対応できるようにしています。

AtCoderのサンプル問題

ZONeエナジープログラミングコンテスト “HELLO SPACE” C問題

C - MAD TEAM

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

N人から3人選んでチームの総合力を最大化する問題です。

is_okの中で探索中の値以上になるか判定しています。

1n = int(input())
2ability = [list(map(int, input().split())) for _ in range(n)]
3
4
5def is_ok(mid):
6    """二分探索中の判定
7
8    Parameters
9    ----------
10    mid : int
11
12    Returns
13    -------
14    resutl : bool
15    """
16    s = set()
17    for a in ability:
18        s.add(sum(1 << i for i in range(5) if a[i] >= mid))
19    for x in s:
20        for y in s:
21            for z in s:
22                if x | y | z == 31:
23                    return True
24    return False
25
26
27def binary_search(ok, ng):
28    """二分探索
29
30    Parameters
31    ----------
32    ok : int
33    ng : int
34
35    Returns
36    -------
37    ok : int
38    """
39    while ng - ok > 1:
40        mid = (ok + ng) // 2
41        if is_ok(mid):
42            ok = mid
43        else:
44            ng = mid
45    return ok
46
47
48ans = binary_search(0, 10**9+1)
49print(ans)

ABC146 C問題

C - Buy an Integer

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

所持金X円の時に、買うことができる最も大きい整数を求める問題です。

シンプルな二分探索の問題です。

1a, b, x = map(int, input().split())
2
3
4def is_ok(mid):
5    """二分探索中の判定
6
7    Parameters
8    ----------
9    mid : int
10
11    Returns
12    -------
13    resutl : bool
14    """
15    if (a * mid) + (b * len(str(mid))) <= x:
16        return True
17    else:
18        return False
19
20
21def binary_search(ok, ng):
22    """二分探索
23
24    Parameters
25    ----------
26    ok : int
27    ng : int
28
29    Returns
30    -------
31    ok : int
32    """
33    while ng - ok > 1:
34        mid = (ok + ng) // 2
35        if is_ok(mid):
36            ok = mid
37        else:
38            ng = mid
39    return ok
40
41
42ans = binary_search(0, 10**9+1)
43print(ans)

ABC192 D問題

D - Base n

AtCoder is a programming contest site for anyone from beginners to experts. We hold weekly programming contests online.

Xをn進数とみなして得られる値で、M以下になるものは何種類あるか求める問題です。

n進数から10進数に変換しながらis_okでM以下か判定します。

1x = list(input())
2m = int(input())
3X = int(''.join(x))
4max_x = int(max(x))
5
6
7def change_nto10(num, b):
8    n = 0
9    numlist = list(str(num))
10    while(numlist):
11        n *= b
12        n += int(numlist.pop(0))
13    return n
14
15
16def is_ok(mid):
17    """二分探索中の判定
18
19    Parameters
20    ----------
21    mid : int
22
23    Returns
24    -------
25    resutl : bool
26    """
27    t = int(change_nto10(X, mid))
28    if t <= m:
29        return True
30    else:
31        return False
32
33
34def binary_search(ok, ng):
35    """二分探索
36
37    Parameters
38    ----------
39    ok : int
40    ng : int
41
42    Returns
43    -------
44    ok : int
45    """
46    while ng - ok > 1:
47        mid = (ok + ng) // 2
48        if is_ok(mid):
49            ok = mid
50        else:
51            ng = mid
52    return ok
53
54
55if len(x) == 1:
56    if X > m:
57        ans = 0
58    else:
59        ans = 1
60else:
61    ans = binary_search(max_x, 10**18+1) - max_x
62
63print(ans)