辞書で単純な条件分岐をシンプルに

if文で条件分岐したのちに何か処理をする...といったことを書くとき、冗長になる時があります。

それをスマートに書く方法が辞書を使う方法です。

LeetCode - 150. Evaluate Reverse Polish Notationを題材にします。

逆ポーランド記法を評価するのに、スタックを用いてやります。

例)

入力として、["4", "13", "5", "/", "+"] が与えられたら、(4 + (13 / 5)) を計算して、6を返します。

この時に冗長な書き方としては、

class Solution:
    def evalRPN(self, tokens: List[str]) -> int:
        stack = []
        for token in tokens:
            if token == "+":
                num1, num2 = stack.pop(), stack.pop()
                stack.append(num2 + num1)
            elif token == "-":
                num1, num2 = stack.pop(), stack.pop()
                stack.append(num2 - num1)
            elif token == "*":
                num1, num2 = stack.pop(), stack.pop()
                stack.append(num2 * num1)
            elif token == "/":
                num1, num2 = stack.pop(), stack.pop()
                stack.append(int(num2 / num1))
            else:
                stack.append(int(token))
        return stack[-1]

となります。冗長、と表現しましたがもしかしたらこちらの方がわかりやすいという方もいるかもしれません。

今回はコード量を減らしたいという一心です笑

辞書を用いると...

class Solution:
    def evalRPN(self, tokens: List[str]) -> int:
        stack = []
        operations = {
            "+": lambda a, b: a + b,
            "-": lambda a, b: a - b,
            "*": lambda a, b: a * b,
            "/": lambda a, b: int(a / b)
        }
        for token in tokens:
            if token in "+-*/":
                operation = operations[token]
                num1, num2 = stack.pop(), stack.pop()
                stack.append(operation(num2, num1))
            else:
                stack.append(int(token))
        return stack[-1]

となります。

スマート！