QuizKnockの鶴崎さんが解いていたパズルのプログラムを自分も書いてみた

QuizKnockのこちらの動画が競プロerの間で話題になっていましたね！
まだ見てない方はご覧ください！

鶴崎さんはクイズ王しながら黄コーダーだったんですよね... やばすぎ
憧れます

パズルの概要はこんな感じです！

3,4,7,8 を使って四則演算で 10を作れみたいなやつ
3,4,7,8 の部分は4個だったり5個だったり6個だったり可変で与えられる
10の部分も入力で与えられる
自由に並び替えて良い、括弧を使って良い
34みたいに桁をくっつけちゃダメ
階乗とかもだめ (四則演算じゃないからね)
単項マイナス (3を-3にするとか) も多分だめ (言及されてないけど)

難しいところ

動画内でも説明されていますが、、、
まず素朴な解法を考えると、順列全探索という方法が浮かびます
並び替えた順列を全て試し、左から全ての +-*/ の組み合わせを適用して試していく、というものです
一見良さそうに見えるのですが、この方法ではだめなケースがあり、
3*4+5*6 のように、トーナメントのような順序で計算されるパターンがカバーできません
そこで、「式」という概念を以下のように考えました
「式」は、以下のどちらかからなる

単項 (1 など)
左の式、演算子、右の式の組み合わせ (式 + 式など)

これを使うと、括弧を自由に使って計算順序を指定したときの二分木の構造が表せます！
(後で教えていただいたんですが、この部分は逆ポーランド記法 を使うとスタックで実装できて楽だったかもしれません)

また、途中の計算結果は、bitDPみたいな持ち方をしました
例えば、コード中の results[13][5] の部分には、A[0], A[2], A[3] を全て使って 5 という結果になる式があれば、それが入ります
(※13 は 2進数では 1101 で、0番目と2番目と3番目のbitが立っているため)

コード

コード公開処刑の時間です
見たい方はどうぞ
GitHubはこちら

クリックすると展開されます

class Formula:
    def __init__(self, used_bit, x, left_formula=None, right_formula=None, operator=None):
        self.used_bit = used_bit
        self.x = x
        self.left_formula = left_formula
        self.right_formula = right_formula
        self.operator = operator
        self.has_multi_terms = operator is not None and operator in '+-'
    def add(self, right_formula):
        assert (self.used_bit & right_formula.used_bit) == 0
        return Formula(self.used_bit|right_formula.used_bit, self.x + right_formula.x, self, right_formula, '+')
    def sub(self, right_formula):
        assert (self.used_bit & right_formula.used_bit) == 0
        return Formula(self.used_bit|right_formula.used_bit, self.x - right_formula.x, self, right_formula, '-')
    def prod(self, right_formula):
        assert (self.used_bit & right_formula.used_bit) == 0
        return Formula(self.used_bit|right_formula.used_bit, self.x * right_formula.x, self, right_formula, '*')
    def div(self, right_formula):
        assert (self.used_bit & right_formula.used_bit) == 0
        assert right_formula.x != 0
        return Formula(self.used_bit|right_formula.used_bit, self.x / right_formula.x, self, right_formula, '/')
    def __str__(self):
        if self.left_formula is None:
            return str(self.x)
        # 余計な括弧はつけないぞという謎のこだわりポイント
        _left = str(self.left_formula)
        left = '({})'.format(_left) if self.left_formula.has_multi_terms and self.operator in '*/' else _left
        _right = str(self.right_formula)
        if self.right_formula.has_multi_terms:
            right = '({})'.format(_right) if self.right_formula.has_multi_terms and self.operator != '+' else _right
        else:
            right = '({})'.format(_right) if self.right_formula.has_multi_terms and self.operator == '/' \
                and self.right_formula.left_formula is not None \
                or self.right_formula.operator == '/' else _right
        return '{}{}{}'.format(left, self.operator, right)

from fractions import Fraction
class MakeXSolver:
    def __init__(self, A, X):
        assert len(A) <= 10, 'too many elements to solve: {}'.format(len(A))
        self.A = A
        self.X = X
        self.N = len(A)
    def solve(self):
        # results[s][x] には、集合sに含まれる要素i番目のA[i]をすべて使ってxを作る作り方を高々1つ格納していく
        results = [{} for _ in range(1<<self.N)]
        for i,a in enumerate(self.A):
            x = Fraction(a)
            results[1<<i][x] = Formula(1<<i, x)
        for s in range(1,1<<self.N):
            if bin(s).count('1') <= 1: continue
            #sの部分集合tを列挙し、u=s\t との組み合わせで作れる結果をすべて results[s]に入れる
            t = s
            while t:
                t = (t-1)&s
                u = s^t
                assert (t&u)==0
                for x,l in results[t].items():
                    for y,r in results[u].items():
                        if x+y not in results[s] and t < u:
                            results[s][x+y] = l.add(r)
                        if x-y not in results[s]:
                            results[s][x-y] = l.sub(r)
                        if x*y not in results[s] and t < u:
                            results[s][x*y] = l.prod(r)
                        if y != 0 and x/y not in results[s]:
                            results[s][x/y] = l.div(r)
        if self.X in results[-1]:
            print('{}={}'.format(results[-1][X], X))
        else:
            print('{} is not found'.format(X))
            print('found results are below')
            for x,f in sorted(results[-1].items()):
                if x.denominator==1:
                    print('{}={}'.format(f, x))
                else:
                    print('{}={}/{}'.format(f, x.numerator, x.denominator))

*A,X = list(map(int,input().split()))
solver = MakeXSolver(A,X)
solver.solve()

実行結果の例

左結合だけじゃなかったり括弧を使ったりしてますが、ちゃんと解を見つけられていそうです

余計な括弧を使わないという謎のこだわりポイント

最初、例えば 3*(4/3+7)=25 という式は (3*((4/3)+7))=25 のように出力していました
が、前者のほうがすっきりしていいですよね！？　前者を出力する改善を地味にやってました
結果、メンテする気の起こらないやばコードが生成されました...
なんだかんだ100個ぐらいは結果を流し見したので大丈夫だと信じたい！
(TODO) ちゃんと頭が回るときに検証したり説明を書いたりするつもり
テストコードないのやばいね

計算量

激ヤバだと思います　集合の部分集合列挙するところ、内側が定数なら $O(3^N)$ らしいですが、そこの内側がやばいですね
部分集合t を全部使ったときにあり得る結果全て * 部分集合u を全部使ったときにあり得る結果全て * 4種類の演算を全部試しています
(TODO) 計算量あとでちゃんと考える
5個までは一瞬、6個だと数秒かかります
7個は実行するのが怖い...