Pull to refresh

Парсинг формул в 50 строк на Python

Reading time 4 min
Views 43K
Вдохновение — задача с собеседования Яндекса и статья «Парсинг формул в 40 строк».

Моей целью было посмотреть, как будет выглядеть «pythonic» решение этой задачи. Хотелось, чтобы решение было простым, код читаемым и разделённым. В итоге ещё получился и пример применения цепочки генераторов (generators pipeline).

На классический алгоритм решения задачи указал Яндекс в своей статье — Алгоритм сортировочной станции.

Алгоритм преобразует выражения в инфиксной, привычной нам, нотации в обратную польскую.
Вычисление выражения в обратной польской нотации (ОПН) для нас привлекательно тем, что у него простой алгоритм.

Весь алгоритм вычисления выражения разбивается на три части:

  1. парсинг исходной строки на числа и операторы
  2. применение алгоритма сортировочной станции для получения выражения в ОПН
  3. вычисление выражения в ОПН

На выходе этапов 1 и 2 мы получаем массивы из чисел и операторов. Велик соблазн реализовать функцию как цепочку генераторов. Мы сократим потребление памяти, используя «ленивую» обработку данных, где выражение вычисляется по мере поступления чисел и операторов.

Итак, приступим


Для начала, определяем операторы в виде словаря, для каждого символа определим приоритет и функцию вычисления.
Этот словарь нам пригодится также для того, чтобы определять, является ли символ оператором:

OPERATORS = {'+': (1, lambda x, y: x + y), '-': (1, lambda x, y: x - y),
             '*': (2, lambda x, y: x * y), '/': (2, lambda x, y: x / y)}

Определим нашу функцию eval_, на вход которой подаётся строка с вычисляемым выражением:

def eval_(formula_string):

Внутри функции определим функцию и два генератора, каждый из которых будет выполнять свою часть работы.

1. Парсер исходной строки

Генератор, получает на вход строку, возвращает числа в формате float, операторы и скобки в формате символов.

    def parse(formula_string):
        number = ''
        for s in formula_string:
            if s in '1234567890.': # если символ - цифра, то собираем число
                number += s  
            elif number: # если символ не цифра, то выдаём собранное число и начинаем собирать заново
                yield float(number) 
                number = ''
            if s in OPERATORS or s in "()": # если символ - оператор или скобка, то выдаём как есть
                yield s 
        if number:  # если в конце строки есть число, выдаём его
            yield float(number)  

2. Алгоритм сортировочной станции

Генератор, получает на вход итерируемый объект из чисел и операторов в инфиксной нотации, возвращает числа и операторов в обратной польской записи.

    def shunting_yard(parsed_formula):
        stack = []  # в качестве стэка используем список
        for token in parsed_formula:
            # если элемент - оператор, то отправляем дальше все операторы из стека, 
            # чей приоритет больше или равен пришедшему,
            # до открывающей скобки или опустошения стека.
            # здесь мы пользуемся тем, что все операторы право-ассоциативны
            if token in OPERATORS: 
                while stack and stack[-1] != "(" and OPERATORS[token][0] <= OPERATORS[stack[-1]][0]:
                    yield stack.pop()
                stack.append(token)
            elif token == ")":
                # если элемент - закрывающая скобка, выдаём все элементы из стека, до открывающей скобки,
                # а открывающую скобку выкидываем из стека.
                while stack:
                    x = stack.pop()
                    if x == "(":
                        break
                    yield x
            elif token == "(":
                # если элемент - открывающая скобка, просто положим её в стек
                stack.append(token)
            else:
                # если элемент - число, отправим его сразу на выход
                yield token
        while stack:
            yield stack.pop()

3. Вычислитель

Функция, получает на вход итерируемый объект чисел и операторов в обратной польской нотации, возвращает результат вычисления:

    def calc(polish):
        stack = []
        for token in polish:
            if token in OPERATORS:  # если приходящий элемент - оператор,
                y, x = stack.pop(), stack.pop()  # забираем 2 числа из стека
                stack.append(OPERATORS[token][1](x, y)) # вычисляем оператор, возвращаем в стек
            else:
                stack.append(token)
        return stack[0] # результат вычисления - единственный элемент в стеке

В конце концов, составляем цепочку генераторов для вычисления результата функции eval_:

    return calc(shunting_yard(parse(formula_string))) 

Быстродействие


Самый главный вопрос: «Как быстро работает программа?» Сравним нашу функцию со встроенной функцией eval.

На простейших случаях наша функция даже быстрее!

%timeit eval("2+2")
100000 loops, best of 3: 12.8 µs per loop

%timeit eval_("2+2")
100000 loops, best of 3: 7.61 µs per loop

На выражениях посложнее — на 22% дольше:

%timeit eval("15/(7-(1+1))*3-(2+(1+1))")
10000 loops, best of 3: 29.7 µs per loop

%timeit eval_("15/(7-(1+1))*3-(2+(1+1))")
10000 loops, best of 3: 36.3 µs per loop

На выражениях ещё сложнее — разрыв увеличивается, но всё равно быстродействие нашей функции сравнимо со встроенной:

%timeit eval("15/(7-(1+1))*3-(2+(1+1))*15/(7-(1+1))*3-(2+(1+1))*(15/(7-(1+1))*3-(2+(1+1))+15/(7-(1+1))*3-(2+(1+1)))")
10000 loops, best of 3: 86.3 µs per loop

%timeit eval_("15/(7-(1+1))*3-(2+(1+1))*15/(7-(1+1))*3-(2+(1+1))*(15/(7-(1+1))*3-(2+(1+1))+15/(7-(1+1))*3-(2+(1+1)))")
10000 loops, best of 3: 147 µs per loop

Да, вспоминая название статьи, тут всего 50 строк, не забывая про читаемость и PEP8!

Код функции целиком
OPERATORS = {'+': (1, lambda x, y: x + y), '-': (1, lambda x, y: x - y),
             '*': (2, lambda x, y: x * y), '/': (2, lambda x, y: x / y)}


def eval_(formula):
    def parse(formula_string):
        number = ''
        for s in formula_string:
            if s in '1234567890.':
                number += s
            elif number:
                yield float(number)
                number = ''
            if s in OPERATORS or s in "()":
                yield s
        if number:
            yield float(number)

    def shunting_yard(parsed_formula):
        stack = []
        for token in parsed_formula:
            if token in OPERATORS:
                while stack and stack[-1] != "(" and OPERATORS[token][0] <= OPERATORS[stack[-1]][0]:
                    yield stack.pop()
                stack.append(token)
            elif token == ")":
                while stack:
                    x = stack.pop()
                    if x == "(":
                        break
                    yield x
            elif token == "(":
                stack.append(token)
            else:
                yield token
        while stack:
            yield stack.pop()

    def calc(polish):
        stack = []
        for token in polish:
            if token in OPERATORS:
                y, x = stack.pop(), stack.pop()
                stack.append(OPERATORS[token][1](x, y))
            else:
                stack.append(token)
        return stack[0]

    return calc(shunting_yard(parse(formula)))

Tags:
Hubs:
+26
Comments 9
Comments Comments 9

Articles