Arantium Maestum

プログラミング、囲碁、読書の話題

部分和問題

06/19/2018追記:コードを書き直した&AtCoderの類題を解いた

プログラミングコンテストチャレンジブック、通称蟻本を読み始めてみる。

とりあえずPythonで解いていって、C++の勉強が進んである程度綺麗に書ける気がしてきたらC++でもやってみたい。

まずは部分和問題。

整数{a_1, a_2, ..., a_n}の部分集合の和が整数kと等しくなるものがあるか判定

1 <= n <= 20

-108 < {a_i} < 108

-108 < {k} < 108

{a_1}から始めて各{a_i}を含むか否かで枝分かれする二分木に対するDFSと見なせる。

def solve(as_, k):
  xs = []
  def dfs(i):
    if sum(xs) == k:
      s=' + '.join(str(x) for x in xs)
      yield "Yes ({k} = {s})".format(k=k, s=s)
    if i < len(as_):
      xs.append(as_[i])
      yield from dfs(i+1)
      xs.pop()
      yield from dfs(i+1)
  for v in dfs(0):
    return v
  return "No"

{a_i}やkの定義域がマイナスも含んでいるので、最悪の場合(和がkになる部分集合が見つからなかったか、探索の最後で見つかったか)220=1048576のケースを全部調べることになる。

今回は解がひとつ見つかればいい、ということなのでdfsが葉ノードに到達しなくてもyieldできるようにした。すべての解を列挙するならdfsの初めの部分を:

def dfs(i):
  if i == len(as_) and sum(xs) == k:
    s=' + '.join(str(x) for x in xs)
    yield "Yes ({k} = {s})".format(k=k, s=s)

こう変えて葉ノードでのみ解がyieldされるようにするのが最も簡単な方法だが、その場合:

yield from dfs(i+1)
xs.append(as_[i])
yield from dfs(i+1)
xs.pop()

と、まずは要素を含まない部分集合から調べ始める方がいいかもしれない。