クルトンのプログラミング教室

Pythonの使い方やPythonを使った競技プログラミングの解法などを解説しています。

Pythonで理解する蟻本「1-6 ハードルが上がった「くじびき」」(p.25)

Pythonで理解する蟻本（プログラミングコンテストチャレンジブック）

この記事は「プログラミングコンテストチャレンジブック第2版」（蟻本）の
「ハードルが上がった「くじびき」」(p.25)
のコードをPythonで書き直したものとなっています。

入力

n\,m\\k_1\,…\,k_n

入力例



3 10
1 3 5


解答

二分探索とO({n}^3\,log\,{n})のアルゴリズム

# 入力
n, m = map(int,input().split())
k = list(map(int,input().split()))

def binary_search(x):
    # xの存在し得る範囲k[l], k[l + 1], ..., k[r - 1]
    l = 0
    r = n
    
    # 範囲に何も含まれなくなるまで繰り返す
    while r - l >= 1:
        i = (l + r) // 2
        if k[i] == x:
            return True # 見つかった
        elif k[i] < x:
            l = i + 1
        else:
            r = i
    # 見つからなかった
    return False

# 二分探索を行うためにソート
k.sort()

f = False

for a in range(n):
    for b in range(n):
        for c in range(n):
            # 最も内側のループの代わりに二分探索
            if binary_search(m - k[a] - k[b] - k[c]):
                f = True

if f:
    print("Yes")
else:
    print("No")

O({n}^2\,log\,{n})のアルゴリズム

# 入力
n, m = map(int,input().split())
k = list(map(int,input().split()))

# 2つで作れる数を格納する配列
kk = [0] * (n * n)

def binary_search(x):
    # xの存在し得る範囲kk[l], kk[l + 1], ..., kk[r - 1]
    l = 0
    r = n * n
    
    # 範囲に何も含まれなくなるまで繰り返す
    while r - l >= 1:
        i = (l + r) // 2
        if kk[i] == x:
            return True # 見つかった
        elif kk[i] < x:
            l = i + 1
        else:
            r = i
    # 見つからなかった
    return False

# k[c] + k[d]の取り得る数を列挙
for c in range(n):
    for d in range(n):
        kk[c * n + d] = k[c] + k[d]

# 二分探索を行うためにソート
kk.sort()

f = False
for a in range(n):
    for b in range(n):
        # 内側の2つのループの代わりに二分探索
        if binary_search(m - k[a] - k[b]):
            f = True

if f:
    print("Yes")
else:
    print("No")

モジュールを使った解法

27ページに書かれているように、このbinary_search関数のような関数は標準で使うことができます。

Pythonではbisectモジュールを使うことで上記と同様の処理を行うことができます。

bisectモジュールを使ったサンプルコードを下に載せておきます。

二分探索とO({n}^3\,log\,{n})のアルゴリズム

# 二分探索するためにbisect_leftをimportする
from bisect import bisect_left

# 入力
n, m = map(int,input().split())
k = list(map(int,input().split()))

# 二分探索を行うためにソート
k.sort()

f = False

for a in range(n):
    for b in range(n):
        for c in range(n):
            # 最も内側のループの代わりに二分探索
            i = bisect_left(k, m - k[a] - k[b] - k[c])
            if 0 <= i < n and k[i] == m - k[a] - k[b] - k[c]:
                f = True

if f:
    print("Yes")
else:
    print("No")

O({n}^2\,log\,{n})のアルゴリズム

# 二分探索するためにbisect_leftをimportする
from bisect import bisect_left

# 入力
n, m = map(int,input().split())
k = list(map(int,input().split()))

# 2つで作れる数を格納する配列
kk = [0] * (n * n)

# k[c] + k[d]の取り得る数を列挙
for c in range(n):
    for d in range(n):
        kk[c * n + d] = k[c] + k[d]

# 二分探索を行うためにソート
kk.sort()

f = False
for a in range(n):
    for b in range(n):
        # 内側の2つのループの代わりに二分探索
        i = bisect_left(kk, m - k[a] - k[b])
        if 0 <= i < n * n and kk[i] == m - k[a] - k[b]:
            f = True

if f:
    print("Yes")
else:
    print("No")


bisectモジュールの使い方を勉強する際には、bisect_leftとbisect_rightの違いを理解し、その使い分けを意識して勉強するとよいでしょう。