Algorithm Search⌘K

Shunting-yard algorithm in Python

Converts infix expressions to postfix or builds expression trees using an operator stack.

File
shunting_yard_algorithm.py, 87 lines
Level
I0 not checked yet
Toolchain
Python 3.14, pytest 9
Test framework
pytest
Run the tests
python3 -m pytest
shunting_yard_algorithm.pyRawGitHub
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/shunting-yard-algorithm/. The authors are listed in prov/2TMV07.jsonl.

"""Dijkstra's shunting-yard algorithm: infix tokens to postfix tokens.

A token is a string. "(" and ")" are parentheses, a key of the operator
table is a binary operator, and any other token is an operand. The operator
table maps each operator to (precedence, associativity), where a larger
precedence binds tighter and associativity is "left" or "right". All
operators of one precedence level must have the same associativity.
"""

from typing import Mapping, Sequence

DEFAULT_OPS: dict[str, tuple[int, str]] = {
    "+": (1, "left"),
    "-": (1, "left"),
    "*": (2, "left"),
    "/": (2, "left"),
    "^": (3, "right"),
}


def _pops(top: tuple[int, str], t: tuple[int, str]) -> bool:
    """True when the operator on the stack must be output before t is pushed."""
    return top[0] > t[0] or (top[0] == t[0] and t[1] == "left")


def shunting_yard(tokens: Sequence[str], ops: Mapping[str, tuple[int, str]] = DEFAULT_OPS) -> list[str]:
    """Return the postfix form of the infix expression `tokens`.

    Raises ValueError when the input is not a well-formed infix expression
    (unbalanced parentheses, two operands or two operators in a row, an
    operator at the start or the end, or an empty expression).
    Time and extra space Theta(n) for n tokens: each token is pushed and
    popped at most once.
    """
    out: list[str] = []
    stack: list[str] = []  # @line L1
    expect_operand = True  # the input check: operands and operators alternate
    for t in tokens:  # @line L2
        if t == "(":
            if not expect_operand:
                raise ValueError("an operand or '(' cannot follow an operand")
            stack.append(t)  # @line L4
        elif t == ")":  # @line L5
            if expect_operand:
                raise ValueError("')' must follow an operand")
            while stack and stack[-1] != "(":  # @line L6
                out.append(stack.pop())
            if not stack:
                raise ValueError("unbalanced ')'")
            stack.pop()  # @line L7
        elif t in ops:  # @line L8
            if expect_operand:
                raise ValueError(f"operator {t!r} needs a left operand")
            p = ops[t]
            while stack and stack[-1] != "(" and _pops(ops[stack[-1]], p):  # @line L9
                out.append(stack.pop())  # @line L10
            stack.append(t)  # @line L11
            expect_operand = True
        else:
            if not expect_operand:
                raise ValueError("two operands in a row")
            out.append(t)  # @line L3
            expect_operand = False
    if expect_operand:
        raise ValueError("the expression is empty or ends with an operator")
    while stack:  # @line L12
        if stack[-1] == "(":  # @line L13
            raise ValueError("unbalanced '('")
        out.append(stack.pop())  # @line L14
    return out  # @line L15


def evaluate_postfix(postfix: Sequence[str], ops: Mapping[str, tuple[int, str]] = DEFAULT_OPS) -> object:
    """Build the expression tree of a postfix sequence as nested tuples (op, left, right)."""
    st: list[object] = []
    for t in postfix:
        if t in ops:
            b, a = st.pop(), st.pop()
            st.append((t, a, b))
        else:
            st.append(t)
    if len(st) != 1:
        raise ValueError("not a postfix expression")
    return st[0]
test_shunting_yard_algorithm.py 185 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/shunting-yard-algorithm/. The authors are listed in prov/2TMV07.jsonl.

import itertools
import random

import pytest

from shunting_yard_algorithm import evaluate_postfix, shunting_yard

# The unambiguous precedence grammar that defines the meaning of an expression.
GRAMMAR = [
    ("E", ("E", "+", "T")), ("E", ("E", "-", "T")), ("E", ("T",)),
    ("T", ("T", "*", "F")), ("T", ("T", "/", "F")), ("T", ("F",)),
    ("F", ("P", "^", "F")), ("F", ("P",)),
    ("P", ("(", "E", ")")), ("P", ("a",)), ("P", ("b",)),
]
NTS = {"E", "T", "F", "P"}


def spans(w):
    """Least fixpoint of the facts 'A derives w[i:j]' (an independent oracle)."""
    n = len(w)
    known = set()

    def ends_of(rhs, i, j):
        ends = {i}
        for x in rhs:
            nxt = set()
            for p in ends:
                if x in NTS:
                    nxt |= {q for q in range(p, j + 1) if (x, p, q) in known}
                elif p < j and w[p] == x:
                    nxt.add(p + 1)
            ends = nxt
        return ends

    changed = True
    while changed:
        changed = False
        for a, rhs in GRAMMAR:
            for i in range(n + 1):
                for j in range(i, n + 1):
                    if (a, i, j) not in known and j in ends_of(rhs, i, j):
                        known.add((a, i, j))
                        changed = True
    return known


def grammar_tree(w):
    """The unique parse tree of w, as nested (op, left, right) tuples, or None."""
    known = spans(w)
    n = len(w)
    if ("E", 0, n) not in known:
        return None

    def build(a, i, j):
        for lhs, rhs in GRAMMAR:
            if lhs != a:
                continue
            for cut in itertools.product(range(i, j + 1), repeat=len(rhs) - 1):
                bounds = (i,) + cut + (j,)
                if any(bounds[k] > bounds[k + 1] for k in range(len(rhs))):
                    continue
                ok = True
                for x, p, q in zip(rhs, bounds, bounds[1:]):
                    if x in NTS:
                        ok = ok and (x, p, q) in known
                    else:
                        ok = ok and q == p + 1 and w[p] == x
                if not ok:
                    continue
                if len(rhs) == 3 and rhs[1] in "+-*/^":
                    return (rhs[1], build(rhs[0], bounds[0], bounds[1]), build(rhs[2], bounds[2], bounds[3]))
                if rhs[0] == "(":
                    return build("E", bounds[1], bounds[2])
                if rhs[0] in NTS:
                    return build(rhs[0], i, j)
                return w[i]
        raise AssertionError("no split")

    return build("E", 0, n)


def test_example() -> None:
    toks = "a + b * ( c - d ) ^ e ^ f - g".split()
    assert shunting_yard(toks) == "a b c d - e f ^ ^ * + g -".split()


def test_associativity_and_precedence() -> None:
    assert shunting_yard("a - b - c".split()) == "a b - c -".split()
    assert shunting_yard("a ^ b ^ c".split()) == "a b c ^ ^".split()
    assert shunting_yard("a + b * c".split()) == "a b c * +".split()
    assert shunting_yard("( a + b ) * c".split()) == "a b + c *".split()
    assert shunting_yard(["x"]) == ["x"]
    assert shunting_yard("( ( x ) )".split()) == ["x"]


@pytest.mark.parametrize("bad", ["", "+", "a +", "+ a", "a b", "( a", "a )", "( )", ") a (", "a ( b )", "a + * b", "( a + ) b"])
def test_rejects_malformed(bad: str) -> None:
    with pytest.raises(ValueError):
        shunting_yard(bad.split())


def random_tree(rng, depth):
    if depth == 0 or rng.random() < 0.3:
        return rng.choice("abcdxyz")
    return (rng.choice("+-*/^"), random_tree(rng, depth - 1), random_tree(rng, depth - 1))


PREC = {"+": 1, "-": 1, "*": 2, "/": 2, "^": 3}


def infix(tree, rng):
    """Print a tree with the parentheses it needs, plus some extra ones."""
    if isinstance(tree, str):
        s = [tree]
    else:
        op, l, r = tree
        ls, rs = infix(l, rng), infix(r, rng)
        lneed = not isinstance(l, str) and (PREC[l[0]] < PREC[op] or (PREC[l[0]] == PREC[op] and op == "^"))
        rneed = not isinstance(r, str) and (PREC[r[0]] < PREC[op] or (PREC[r[0]] == PREC[op] and op != "^"))
        s = (["("] + ls + [")"] if lneed else ls) + [op] + (["("] + rs + [")"] if rneed else rs)
    if rng.random() < 0.1:
        s = ["("] + s + [")"]
    return s


def postorder(tree):
    out, stack = [], [(tree, False)]
    while stack:
        t, done = stack.pop()
        if isinstance(t, str):
            out.append(t)
        elif done:
            out.append(t[0])
        else:
            stack.append((t, True))
            stack.append((t[2], False))
            stack.append((t[1], False))
    return out


def test_random_trees() -> None:
    rng = random.Random(1961)
    for _ in range(1500):
        tree = random_tree(rng, rng.randint(0, 5))
        toks = infix(tree, rng)
        post = shunting_yard(toks)
        assert post == postorder(tree)
        assert evaluate_postfix(post) == tree


def test_random_token_strings_against_grammar() -> None:
    rng = random.Random(42)
    alphabet = ["a", "b", "+", "-", "*", "^", "(", ")"]
    for _ in range(1500):
        w = [rng.choice(alphabet) for _ in range(rng.randint(0, 9))]
        expected = grammar_tree(w)
        if expected is None:
            with pytest.raises(ValueError):
                shunting_yard(w)
        else:
            assert evaluate_postfix(shunting_yard(w)) == expected


def test_all_short_strings() -> None:
    alphabet = ["a", "-", "^", "(", ")"]
    for n in range(0, 6):
        for w in itertools.product(alphabet, repeat=n):
            w = list(w)
            expected = grammar_tree(w)
            if expected is None:
                with pytest.raises(ValueError):
                    shunting_yard(w)
            else:
                assert evaluate_postfix(shunting_yard(w)) == expected


def test_long_input() -> None:
    n = 20000
    toks = ["("] * n + ["a"] + [")"] * n
    assert shunting_yard(toks) == ["a"]
    toks = ["a"] + ["^", "a"] * n
    assert shunting_yard(toks) == ["a"] * (n + 1) + ["^"] * n

What the next level needs

  • the tests did not run in this build

Read the full page about Shunting-yard algorithm