組み合わせの列挙で学ぶ再帰関数

こんにちは，今回もプログラミングコンテスト攻略のためのアルゴリズムとデータ構造をPythonを用いて実装しながら，解説していきたいと思います．

目次

再帰・分割統治法

問題を解く時に，部分的に問題を解く関数を用意することで与えられた元の問題を解いていくテクニックの一つが分割統治法です．
再帰関数とは，関数の中で自分自身をさらに呼び出すような関数のことを言います．

全探索の問題

今回は全探索の問題を再帰を使って解いていきたいと思います．実際により効率的に解こうと思えば，他の方法もありますが，再帰に慣れるため計算時間のかかる方法で解いています．
問題文は以下の通りです．
長さnの数列Aと整数mに対して，配列Aの要素の中のいくつかの要素を足し合わせてmが作る事ができるかどうかを判定するプログラムを作成してください．Aの各要素は一度だけ使う事ができます．

入力：一行目にn,二行目にAを表すn個の整数,
　　　三行目にq，四行目にq個の整数m_iが与えられます.
出力：整数mを作る事ができれば"yes",できなければ”no”
制約：n ≦ 20
　　　q ≦ 200
　　　1 ≦ Aの要素 ≦ 2000
　　　1 ≦ 整数m ≦ 2000

対応する問題はALDS1_5_A:Exhaustive Searchです．

組み合わせ列挙の関数

今回は全探索で判定を行うようなプログラムを作成していきたいと考えているので，その要素を使うか使わないかを示した全組み合わせが必要になります．
まずが，要素の組み合わせを列挙するような関数をまず作ってみたいと思います．

#n=5の時の組み合わせ
def comb(n):
    A = [0]*n
    select(0,A)
        
def select(i, A):
    if i == n:
        print(A)
        return
    select(i+1, A)
    A[i] = 1 #A[i]を選択
    select(i+1, A)
    A[i] =0  #A[i]を選択しない
n = 5
comb(n)

これで組み合わせの列挙はできるようになりました．
また，要素の組み合わせをpythonで作る場合より効率的な方法がいくつかあるので別記事で紹介しています．

tech-shelf.hatenablog.com

全探索を解くためのプログラム

次にこの組み合わせの中で整数mを作ることができるかどうかチェックするプログラムを作成します．
これは先ほどのプログラムのA[i]を選択するかしないかの部分を書き換えていくことで作成することができます．

n = 5
A = [1,5,7,10,21]
q = 4
m = [2,4,17,8]
def solve(i,m):
    if m == 0:
        return 1
    if i >= n:
        return 0
    rec =  solve(i+1,m) or solve(i+1, m - A[i])
    return rec
for i in m:
    if solve(0,i):
        print("yes")
    else:
        print("no")

このプログラムを一般化することでこの問題は解くことができます．

問題の解答

では，問題に答えられるように形を整えていきます．

n = int(input())
A = list(map(int,input().split()))
q = int(input())
m = list(map(int,input().split()))
sumA = sum(A)
def solve(i,m):
    if m == 0:
        return 1
    if i >= n:
        return 0
    rec =  solve(i+1,m) or solve(i+1, m - A[i])
    return rec
for i in m:
    if i >= sumA:
        print("no")
    else:
        if solve(0,i):
            print("yes")
        else:
            print("no")

※Pythonで問題を解く場合，整数mが配列Aの総和より大きい場合再帰関数に飛ばさないという分岐を作成する必要があります．

このアルゴリズムは選択するかしないかというのを各配列の要素に聞くような再帰関数を用いています．そのため，計算量はO（2ⁿ）となっており，nが小さい場合にしか活用することができません．
今後さらに章を進めていく中でもっと効率的なアルゴリズムがでてくるので，まずは全列挙が簡単に作れるようにマスターしましょう．