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

はじめに

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

bisectの使い方

bisectは、ソートされたリストにソートされた状態を保ちながら挿入、挿入する場所を求めることができるライブラリです。

二分探索を利用しているため効率よく挿入する場所を見つけることができます。挿入する場所の探索は\(O(\log n)\)で行えます。

挿入自体の操作は\(O(n)\)かかることに注意してください。

挿入できる場所を求める

bisect_leftは、挿入できるリストの添字を返します。同じ値がある場合は、その値の最も左側の添字になります。

li = [2, 5, 8, 13, 13, 18, 25, 30]

ind = bisect.bisect_left(li, 10)
print(ind)

ind = bisect.bisect_left(li, 13)
print(ind)

以下のようにソートされた状態を保ちながら挿入できる添字を返します。

3
3

bisect_rightbisectは、挿入できるリストの添字を返します。同じ値がある場合は、その値の最も右側の添字になります。

li = [2, 5, 8, 13, 13, 18, 25, 30]

ind = bisect.bisect_right(li, 10)
print(ind)

ind = bisect.bisect_right(li, 13)
print(ind)

ind = bisect.bisect(li, 10)
print(ind)

ind = bisect.bisect(li, 13)
print(ind)

以下のようにソートされた状態を保ちながら挿入できる添字を返します。

3
5
3
5

挿入する

bisectではリストに挿入もできます。ただし、挿入自体の操作は\(O(n)\)かかります。

insort_leftはリストに値を挿入します。同じ値がある場合は最も左側に挿入します。

li = [2, 5, 8, 13, 13, 18, 25, 30]
bisect.insort_left(li, 13)
print(li)
[2, 5, 8, 13, 13, 13, 18, 25, 30]

insort_rightinsortはリストに値を挿入します。同じ値がある場合は最も右側に挿入します。

li = [2, 5, 8, 13, 13, 18, 25, 30]
bisect.insort_right(li, 13)
print(li)
[2, 5, 8, 13, 13, 13, 18, 25, 30]

insortも同じ動きをします。

li = [2, 5, 8, 13, 13, 18, 25, 30]
bisect.insort(li, 13)
print(li)
[2, 5, 8, 13, 13, 13, 18, 25, 30]

AtCoderのサンプル問題

ABC130 D問題

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

長さNの数列Aと整数Kが与えられ、Aの連続する部分列でその要素の和がK以上となるものはいくつあるのか求める問題です。

累積和を求めてから、bisectで挿入できる場所を探索します。

import bisect
import itertools

n, k = map(int, input().split())
a = list(map(int, input().split()))
X = [0] + list(itertools.accumulate(a))
ans = 0

for x in X:
    ind = bisect.bisect_left(X, k+x)
    if ind <= n:
        ans += n - ind + 1

print(ans)

ABC205 D問題

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

長さNの数列AとQ個の\(K_{i}\)が与えられ、数列A含まれない\(K_{i}\)番目の整数を求める問題です。

数列のそれぞれの値以下で選べる整数の個数のリストを作成して、bisectで挿入できる場所から\(K_{i}\)番目の整数を求めます。

import bisect

n, q = map(int, input().split())
A = list(map(int, input().split()))
K = []
for _ in range(q):
    K.append(int(input()))

A.sort()

B = [A[0]-1]
for i in range(1, len(A)):
    B.append(A[i]-A[i-1]+B[i-1]-1)

for k in K:
    if B[-1] < k:
        print(A[-1] + k - B[-1])
    else:
        ind = bisect.bisect_left(B, k)
        print(A[ind] - 1 - (B[ind] - k))

参考

タイトルとURLをコピーしました