再帰関数でpythonのyieldを使う（例題：数字和がnになる数を昇順でリスト）

再帰関数でpythonのyieldをどう使うか分かりづらかったので調べてみました。

【例題】数字和がnになる数を小さい順に出力する関数を作れ

このような問題を例として考えてみます。まず数字和がnになるk桁の数を昇順にプリントする再帰プログラムを考えます。

数字和がnになるk桁の数を昇順にプリントする

def ds2n(n,k,ans=0): # n:数字和, k:桁数, ans:答え 初期値0, 
    if n==0 and k==0:
      print(ans)
      return
    for d in range(max(0 if ans>0 else 1,n-(k-1)*9),10):
      ans1, n1 = ans*10+d, n-d
      if n1 >= 0:
        ds2n(n1,k-1,ans1)
    return

n = 5
for k in range(1,2+1):
  print(f"-- k={k} --")
  ds2n(n,k)
# k=1 --
# 5
# -- k=2 --
# 14
# 23
# 32
# 41
# 50

例題はこのkを1からカウントアップすればよいので関数digsum2nをかぶせます。中身は同じなので省略しています。

数字和がnになる数を昇順にプリントする

def digsum2n(n):     # n:数字和
    :
  for k in range(1,2+1):
    ds2n(n,k)
  return 
digsum2n(5)
#結果は同じ

結果をgeneratorとして出すようにyieldに変える

ここからyieldをどう入れるかですか、基本方針に以下の通りです。

1. Print文をyieldに変える
2. 関数を呼ぶ前にyield fromを入れる

【解答】yieldを使ったプログラム

この方針で簡単にGenaratorとなるプログラムに変更できました。ほぼ無限に数を生成するので呼び側で値を見て止める必要があります。

この考え方は、この問題のように再帰関数などの深いところで答えが見つかる場合などの有効だと思います。

def digsum2n(n):       # n:数字和
  def ds2n(n,k,ans=0): # n:数字和, k:桁数, ans:答え 初期値0, 
    if n==0 and k==0:
      yield ans
      return
    for d in range(max(0 if ans>0 else 1,n-(k-1)*9),10):
      ans1, n1 = ans*10+d, n-d
      if n1 >= 0:
        yield from ds2n(n1,k-1,ans1)
    return

  for k in range(1,100+1):
    yield from ds2n(n,k)
  return 

for ans in digsum2n(20):
  if ans >= 500: break
  print(ans)
# 299
# 389
# 398
# 479
# 488
# 497

（開発環境：Google Colab）

この考え方はProject Euler Problem 254: Sums of Digit Factorialsを解くのに役に立ちます。