Pythonで二分探索法

AtCoderのRegular Contest084に参加した。0完。。気づいてなかったのだが、AtCoderの提出物一覧から、他の人の正答が見れるようだ。

C — Snuke Festival

C問題がどうやっても解けず困っていたのだが、他の方の回答を見ると、Pythonで bisect というライブラリを用いている人が多かった。これは二分探索法を行うモジュールのようだ。

二分探索法は、下記のようなアルゴリズム。

• ある条件を満たすものの最大値を求めるというタスク
• 条件を満たす小さい数aと条件を満たさない大きな数bを考える
• その中間の数cが条件を満たすか確認する
• cが条件を満たすのであればcをaに代入し、再度aとbの中間の数が条件を満たすか確認
• cが条件を満たさないのであればcをbに代入し、再度aとbの中間の数が条件を満たすか確認

こうすることで、探索時間を減らすことができる。これを用いることでTLEにならないようだ。

AtCoderをPythonで解いてる方が少ないのか、CのPythonでの正答者は重数人程度だったが、半数以上が bisect を用いているようだった。

bisect でC問題を解いてみたい。

11/5 追記

ARC 84のC問題の回答をみていたところ、こちらの方の回答がわかりやすかった。

Submission #1747234 - AtCoder Regular Contest 084

import bisect

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 = sorted(a)
b = sorted(b)
c = sorted(c)

ans = 0

for x in b:
    aa = bisect.bisect_left(a,x)
    cc = len(c) - bisect.bisect_right(c,x)
    ans += aa * cc

print(ans)

bisect_left は、第一引数のソートされたリストの何番目に、第二引数が入るか、同じ数字の場合は左側に入れる、とした場合の、位置を返してくれる。

bisect_right は同じ数字の場合は右側に入れた際の位置を返してくれる。