QuizKnockのパズル解きプログラムを実装してみた！

TL;DR

最近Atcoderを始めました！
楽しいですが初参加したAtcoder Beginer ContestではA,B問題しか解けず悔しい思いをしています。

そんな時にQuizknockの動画で鶴崎さんがプログラムでパズル王ふくらpに挑戦するのを拝見しました。

【大決戦】コンピュータを使ってでもパズル王に勝ちたい！！

そのプログラムが練習に良さそうだったので実装してみました。
問題を紹介すると以下の数字パズルです

問題

[n1, n2, ... , n_k]のk個の数字に「＋、ー、×、÷」を適用して目的の数Tを作る。
その数式を文字列として出力せよ

問題例¹

[2, 5, 5, 9]を使って13を作る

答えの１例（答えが複数個存在する可能性あり）

(9 - 5) × 2 + 5

プログラム

さっそくプログラムを見せます。

プログラムではExpressionという名前で答えの式の計算結果と文字列を管理していて、再帰関数で答えのExpressionを求めています。
また数値計算は計算誤差が無いようにすべてfractionsの有理数演算に任せました。

from dataclasses import dataclass
from typing import List

from itertools import combinations
from fractions import Fraction
import copy

@dataclass
class Expression():
    v: Fraction  # valueだけど計算する時、プログラムを書くのが楽になるように一文字にした。
    c: str  # character

# target = 25
# nums = [1,2,2,3,6,9]
        
target = 13
nums = [1,2,3,3,3,9]

e_list = [Expression(v = Fraction(n), c = str(n)) for n in nums]

def get_ans_expression(e_list: List[Expression]) -> Expression:
    length = len(e_list)
    if length == 1:
        return e_list[0]
    indices = range(length)

    # インデックス一覧のindicesから二つ選ぶ
    for i0, i1 in combinations(indices,2):
        e_list_copy = copy.deepcopy(e_list)
        e1 = e_list_copy.pop(i1)
        e0 = e_list_copy.pop(i0)

        # 選んだ二つから新しいExpressionを計算する。
        new_values = [e0.v + e1.v,
                      e0.v - e1.v,
                      e1.v - e0.v,
                      e0.v * e1.v]
        
        new_chars = ["(" + e0.c + "+" + e1.c + ")",
                     "(" + e0.c + "-" + e1.c + ")",
                     "(" + e1.c + "-" + e0.c + ")",
                     "(" + e0.c + "*" + e1.c + ")"]
        
        if e1.v != 0:
            new_values.append(e0.v / e1.v)
            new_chars.append("(" + e0.c + "/" + e1.c + ")")
        if e0.v != 0:
            new_values.append(e1.v / e0.v)
            new_chars.append("(" + e1.c + "/" + e0.c + ")")

        for new_value, new_char in zip(new_values, new_chars):
            # 2重のforループで、
            # 6通りの演算とn個の中から2個選ぶすべての組み合わせに対して再帰呼び出しをする。
            # 再帰呼び出し先ではn-1個のExpressionで同じことする。
            new_expression = Expression(new_value, new_char)
            r = get_ans_expression(e_list_copy + [new_expression])
            
            # 得られたExpressionが目標と一致していたら、
            # すべての再帰呼び出しを一つずつreturnで抜け出す。
            if r.v == target:
                return r
    
    # 見つからなかった場合
    return Expression(target-Fraction(1),"0")

print(get_ans_expression(e_list))
# -> Expression(v=Fraction(13, 1), c='((3*3)+((3+9)/(1+2)))')

上手く動いてる？

３例確認した範囲ではあっていそうです！

エッジケースなどで間違いがないか少し心配です。
Atcoderと違って、ACのチェックも自分で頑張らないといけませんね。

間違いや改善点などあったら助かります(^^)/

1. テストケースを含めて3例を動画より引用しており、Quizknock作問です。 ↩