【AtCoder】Pythonで二分探索
はじめに
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)
参考
- 解説 - ZONeエナジー プログラミングコンテスト “HELLO SPACE”
- 解説 - AtCoder Beginner Contest 146
- 解説 - SOMPO HD プログラミングコンテスト2021(AtCoder Beginner Contest 192)