Python 競技プログラミング 処理高速化 二分探索

Question

Python 競技プログラミング 処理高速化 二分探索

環境

最近，趣味で競技プログラミングを始めた者です．個人的に使用経験のあるPythonを使っています．
今回のコードの実行環境はatcoderのPython(3.8.2)です．
私がコードの下書きに使っているのは，Google Colabになります．
実行環境でPyPy3(7.3.0)も選べるので利用してみましたが，実行時間はほぼ変わりませんでした．

問題

atcoder内のアルゴリズムと数学 演習問題集の，009 - Brute Force 2の問題になります．
もしかしたら，著作権的な問題があるかもしれませんので，お手数をおかけしますが問題は写しません．

リンク先より確認いただければ，幸いです．

質問

Pythonの処理を高速化する方法として，

1. 標準入力においてinput()ではなく，sys.stdin.readlineを使用する
2. 全探索ではなく二分探索を行う
3. リストlist()ではなく，集合set()を使う

この3点がウェブを調べていてわかりました．1と２の高速化手段を組み合わせたコードと，１と３高速化手段を組み合わせたコードを実装できましたが，満点の条件である実行時間で1000msが切れませんでした．どちらのコードでも1100ms強の実行時間がかかってしまいました．

現状，私ではこれ以上の高速化できないため，他の高速化手段について知りたいと思い，質問しました．

また，Pythonでは，ほぼこれ以上速くならないのであれば，仕方ないので先の問題に進もうと思います．
ちなみに，この問題の元の書籍を購入して読みましたが，基本のコードがC++であるため，参考にはなりませんでした．また，この本の付録的なGithubの参考コードのPythonには，この問題のコードは存在しないようです．

お力添えのほど，よろしくお願いいたします．

1と２の高速化手段を組み合わせたコード

import itertools
import sys
import bisect
input = sys.stdin.readline
a, b = map(int, input().split())
c = input().split()
c = [int(t) for t in c]
s = []
for n in range(a,0,-1):
  for combi in itertools.combinations(c, n):
    s.append(sum(combi))
s.sort()
n = bisect.bisect_left(s, b)
bisect.insort_left(s,b)
if s[n] == s[n+1]:
  print("Yes")
else:
  print("No")

１と３高速化手段を組み合わせたコード

import itertools
import sys
input = sys.stdin.readline
a, b = map(int, input().split())
c = input().split()
c = [int(t) for t in c]
s = set()
for n in range(a,0,-1):
  for combi in itertools.combinations(c, n):
    s.add(sum(combi))
if b in s:
  print("Yes")
else:
  print("No")

ベストアンサーと解決済みコード

本業が忙しく，また動的計画法というものを知らなかったので，解答の理解に時間がかかってしまいました．
kazuma-sさん，lehshellさん，actorbugさん，ご回答いただきありがとうございました．
今回は解決だけでなく，学習につながる解説をしていただいたactorbugさんをベストアンサーとさせていただきます．
actorbugさんが投稿していただいたコードを解釈してみましたので，掲載しておきます．実行時間(63msほど)でした．
私としては，納得できましたが何か間違い等にお気づきの方はコメントいただけると幸いです．

解決済みコードとその解釈

a, b = map(int, input().split()) # aがカードの枚数， bが判定値（判定基準）
c = map(int, input().split()) # cがカードの数値リスト
dp = [True] + [False] * b # dpのノードを作る．最初がtrueなのは， 判定する初期位置の設定？
for n in c: # カードの数値を取り出す．
    for i in range(b - n, -1, -1): # 判定値から，リストの数字を引く．そこから， 値を大きい方から-1ずつ取り出す．indexの0まで
      if dp[i]: # dpのリスト内の値は，真偽値なので帰ってくるのは"True" or "False"だから， Trueの場合は
        dp[i + n] = True # Trueは値としては1だから，dpリスト内の1+（取り出した数値n）をTrueにする．　判定する位置の移動？
print("Yes" if dp[-1] else "No") # リストのindexがTrueならその数値をカードの合計で作れる．最後が判定値なので，"Yes"，Falseなら"No"

Answer 1

満点の条件である実行時間で1000msが切れませんでした．

問題の「注意」にも書いてありますが、この問題は2.4節までの知識(全探索)で解いても1000msを切れません。3.7節の知識(動的計画法)で解いて、初めて1000msを切れるようになっています。

どちらのコードでも1100ms強の実行時間がかかってしまいました．

AtCoderでは、時間がかかるコードは途中で強制終了されます。1100msで実行が終わったわけではなく、1100msで強制終了されただけですので、実行時間の数字自体には意味はありません。

全探索ではなく二分探索を行う

ソートしてから1回だけしか二分探索しない場合、普通に検索するよりかえって遅くなるので注意しましょう。(ソートがO(N log N)なので、普通に検索するO(N)より遅くなる)

---

3.7節の知識(動的計画法)で解く場合、以下のようになります。

a, b = map(int, input().split())
c = map(int, input().split())
dp = [True] + [False] * b
for n in c:
    for i in range(b - n, -1, -1):
        if dp[i]:
            dp[i + n] = True
print("Yes" if dp[-1] else "No")

2.4節までの知識で解く場合でも、和をすべてsetなどに集めてから検索するのではなく、その場で和が目的の値か判定したほうが速くなりそうです。

import itertools
import sys
input = sys.stdin.readline
a, b = map(int, input().split())
c = input().split()
c = [int(t) for t in c]
for n in range(a,0,-1):
  for combi in itertools.combinations(c, n):
    if sum(combi) == b:
      print("Yes")
      exit()
print("No")

なお、書籍の3.7.9節に書いてある通り、この問題には「部分和問題」という名前がついています。こちらの名前で検索すれば、さまざまな解説や実装例が見つかると思います。

---

他の方の回答にあるように、メモ化せずに再帰を使った場合、特定の条件で時間がかかるという問題があります(カードの枚数が多め、目標とする合計が大きめ(全カードの合計付近)、答えがNo)。手元の環境だと、カード枚数23枚で1秒を超えました。

23 277
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23

対策されてしまったので、もう少し意地悪な例を挙げておきます。

23 298
2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24

---

私のコードはいろいろ削ぎ落して分かりづらくなっているので、「部分和問題」で検索して見つかるコード・解説で学ぶことをお勧めします。
動的計画法の解説は、表を使わないと厳しいので、ここの回答欄だと限界があります。

無理を承知で簡単に解説すると、以下のようになります。

まず、dpの意味ですが、for n in c:のループを何回回ったかにより変わってきます。
ループをj回回った時点でdp[i]がTrueならば「cの先頭j枚のカードからいくつか選んで、合計がiとなるようにする方法がある」という意味となります。

初期状態(ループを0回回った時点)では、カードを1枚も使えないので、作ることが可能な合計は0しかありません。そのため、dp[0]のみTrueで、それ以外はFalseになるようにdpを初期化しています。

終了状態(ループをa回回った時点)では、すべてのカードを使える場合の結果が求まっています。そのため、dp[b]がTrueなら、合計がbになるようにする方法がある、ということになります。(dpの長さがb+1なので、dp[b]とdp[-1]は同じ)

さて、肝心なループの中身についてですが、外側のループをj-1回回った時点でdp[i]がTrue、つまり先頭j-1枚のカードで合計iにできるなら、j回目のループ、つまり先頭j枚のカードが使える場合に以下のことが言えます。

1. 合計をiにする方法がある(j枚目のカードを使わなければよい)
2. 合計をi + nにする方法がある(j-1枚のカードで合計iにしたうえで、j枚目のカードnを使う)

この2.の場合を反映するのが、一番内側のif文になります。なお、1.については、dpを上書きしているので何もしなくても反映されます。

あとは内側のループが逆順となっている理由ですが、小さい順に処理してしまうと、dp[i+n]=TrueにてTrueにした(=合計をi+nにするのにj枚目のカードを使用済み)値を、同じループ内で再度拾ってしまうので、j枚目のカードが複数回使われてしまうためです。

Answer 2

n, s = map(int, input().split())
a = list(map(int, input().split()))

def f(i, s):
	if s == 0: return True
	if s < 0 or i == n: return False
	if f(i+1, s - a[i]): return True
	return f(i+1, s)
	
print('Yes' if f(0, s) else 'No')

追記
全部の合計より大きい数は最初から No にしないといけませんね。

print('Yes' if s <= sum(a) and f(0, s) else 'No')

Comments

[tmp]

すばらしい。これに動的計画法？ャッシュみたいなもの？を入れれば完成ですね。

Answer 3

スライス演算削除してさらに高速化（本質は kazuma-s さんと同じです）

def comb_sum(nums, val, sz, idx=0, exist=False):
    if exist and val == 0:
        return True
    if val < 0 or idx >= sz:
        return False
    return (comb_sum(nums, val - nums[idx], sz, idx+1, True) or
            comb_sum(nums, val, sz, idx+1, exist))

n, s = map(int, input().split())
nums = [*map(int, input().split())]
ans = comb_sum(nums, s, len(nums))
print('YES' if ans else 'NO')

高速化しました。

def comb_sum(nums, val, exist=False):
    if exist and val == 0:
        return True
    if val < 0 or not nums:
        return False
    return comb_sum(nums[1:], val - nums[0], True) + comb_sum(nums[1:], val, exist)

n, s = map(int, input().split())
nums = [*map(int, input().split())]
ans = comb_sum(nums, s)
print('YES' if ans else 'NO')

以前のコードです。

def comb_sum(nums, val):
    if val <= 0 or not nums:
        return [[]] if val == 0 else []
    else:
        return [[nums[0]] + n for n in 
                comb_sum(nums[1:], val - nums[0])] + comb_sum(nums[1:], val)

n, s = map(int, input().split())
nums = [*map(int, input().split())]
ans = comb_sum(nums, s)
print('YES' if ans else 'NO')