Codingame 1D Spreadsheet をメモ化を使って解く [Python]

はじめに

最近Codingameの練習問題のEasyを解いているのですが、ちょいちょい詰まることがあります。今回は題名の通り1D Spreadsheetという問題なのですが、これが初心者の自分にはなかなか難しくて放置してまた考えてというふうに繰り返していると数日かかってしまいました。
最初は自力でなんとか解こうと思っていたのですが、最後のテスト問題が解けず、コードを弄っているとどの問題もパスできなくなるなどの負のループに陥ってしまいました。ついに諦めてForumを見てみると、新しい発見とMemoization(メモ化)というアルゴリズムを使うといいということがわかりました。初心者なりにメモ化を使ってみたので、コードに対するフィードバックをもらえると嬉しいです。

問題

問題はセル数と計算セルが与えられて、それの答えをアウトプットするというものです。
計算セルはそれぞれVALUE, ADD, SUB, MULTのいずれかの指令と、それに続く2つのバリューarg1,arg2で構成されています。
指令どおりにarg1とarg2足したり引いたりかけたりするだけだったら簡単なのですが、ここに$を使った参照というものも入ってきます。$0で0番目の答えを参照という具合です。
インプットの例は以下のとおりです。

2
VALUE 3 _
ADD $0 4

そしてこのインプットに対して以下を出力しなければいけません。

3
7

解説すると、最初のセルはVALUEで最初のarg1の3が答えになり、次のセルでは$0と4を足すのですが、$0は最初のセルの答えが入るので答えは7になります。

問題はこの参照が厄介で、前のセルの答えだけでなく、あとのセルも参照しなければいけなかったり、全部が参照で構成されていたりします。

なんとなく説明しただけですので、詳しくは以下をご覧ください。投げやりですみません😂
https://www.codingame.com/ide/puzzle/1d-spreadsheet

コード

セル数の上限が100までで、すべてのセルに対して再計算をさせてやるとタイム・アウトしてしまう可能性があります。そこで本コードではメモ化を用いて再計算を防いでやるというアルゴリズムを使います。

インプット

まずは競技プロ定番のインプットで、問題に対応させてやります。
追加でメモ化のために答えを記録しておく辞書を用意します。

#セル数のインプット
n = int(input())

#インプットされるセルとメモ化のための答えの辞書の用意
cells = []
ans = {}

#セルのインプット
for i in range(n):
    cells.append(input().split())

関数

問題処理のための関数です。i番目のセルの答えが出るようにコードを組んでやります。
いちいち再計算が行われるのを防ぐために、答えが出た時点でans辞書に記録して、他のセルで答えの参照が出たときに再参照と再計算をしないようにしてやります。

def val(i):

    #　iセルの中身をわかりやすいように定義
    function = str(cells[i][0])
    arg_1 = str(cells[i][1])
    arg_2 = str(cells[i][2])

    #　辞書内にセルiの答えがあったときに答えを返す
    if i in ans:
        return ans[i]

    #　'SUB $a - $a'のとき０を返す
    if (function == 'SUB') and (arg_1 == arg_2):
        ans[i] = int(0)
        return ans[i]

    #　arg_1が参照だったとき、関数に投げる
    if '$' in arg_1:
        arg_1 = val(int(arg_1.strip('$')))

    #　arg_1が参照だったとき、関数に投げる
    if '$' in arg_2:
        arg_2 = val(int(arg_2.strip('$')))

    #　最後に参照ではなかったときに計算、'ans'辞書に記録、そして答えを返す。
    if ('$' not in str(arg_1)) and ('$' not in str(arg_2)):
        if function == 'VALUE':
            ans[i] = int(arg_1)
            return ans[i]
        if function == 'ADD':
            ans[i] = int(arg_1) + int(arg_2)
            return ans[i]
        if function == 'SUB':
            ans[i] = int(arg_1) - int(arg_2)
            return ans[i]
        if function == 'MULT':
            ans[i] = int(arg_1) * int(arg_2)
            return ans[i]

アウトプット

最後に答えを出力するために関数に順番にセル数を投げてやります。

for i in range(n):
    print(val(i))

まとめ

解けてしまうとなぜこんなのに時間かかってたんだろうと少し恥ずかしい気持ちになりますが、少しでもこの問題を解いていて行き詰まった人の役に立てばなと思います。

あと、コードを人に見てもらうという機会がないので、自分のコードがきれいか汚いかすらわからないのですが、こう書いたほうがいいよとかあったらコメントお願いします！

Appendix

import sys
import math

n = int(input())

cells = []
ans = {}

for i in range(n):
    cells.append(input().split())
    #ans[i] = None


def val(i):
    function = str(cells[i][0])
    arg_1 = str(cells[i][1])
    arg_2 = str(cells[i][2])

    if i in ans:
        return ans[i]

    if (function == 'SUB') and (arg_1 == arg_2):
        ans[i] = int(0)
        return ans[i]

    if '$' in arg_1:
        arg_1 = val(int(arg_1.strip('$')))

    if '$' in arg_2:
        arg_2 = val(int(arg_2.strip('$')))

    if ('$' not in str(arg_1)) and ('$' not in str(arg_2)):
        if function == 'VALUE':
            ans[i] = int(arg_1)
            return ans[i]
        if function == 'ADD':
            ans[i] = int(arg_1) + int(arg_2)
            return ans[i]
        if function == 'SUB':
            ans[i] = int(arg_1) - int(arg_2)
            return ans[i]
        if function == 'MULT':
            ans[i] = int(arg_1) * int(arg_2)
            return ans[i]


for i in range(n):
    print(val(i))