[Python] Simplify a prime factorization like expression with optimal performance

I am working with a situation in which I need to turn an arithmetic string expression (e.g. (3/2) * (3**4 / (20 * 4))**(-1/4)) into the following prime factorization form: $\prod_{p:\text{prime}}p_i^{k_i}$ (in LaTeX format), where $p_i$ are prime numbers and $k_i$ are rational numbers. The result, I want it to be returned as a dict where the keys are prime factors and values are their respective exponents.

The arithmetic expression contains only multiplication, division, and power operations. The exponent of power can be regarded as a constant (i.e. contains no more recursive sub-structures. When it is a fraction it will be turned to a Fraction object).

Below is my current attempt to solve this problem by traversing the abstract syntax tree recursively. For a leaf node, which is a single integer value, I take its prime factors and store it in a dict; for power I multiply each element in the dict by the exponent and for multiplication and division I add or subtract each term in the dict. This method works, however, it creates a separate dict for every leaf node. As dict is a mutable type, creating it multiple times isn't efficient.

I am looking for a solution which can create only one dict at the beginning and perform all updates in this dict without having to create other dict. Also in my actual use case the situation may be that I already have a dict containing a prime factorization and I want to update the dict so that it represents the multiplication of the dict's factorization and my expression. For the solution of this problem I am investigating into prefix / postfix (Polish / reverse Polish) notation, but I am not sure whether this is the correct direction. So is my requirement of using only one dict (an external data structure, such as a stack is acceptable) possible?

import ast
from fractions import Fraction as Q
from numbers import Integral, Rational
from collections import Counter
from collections.abc import MutableMapping


def _mul(m: MutableMapping[int, Rational], n: Rational):
if n == 0:
m.clear()
return
if n == 1:
return
for k in m:
m[k] = m[k] * n


def _cleanZeroes(m: MutableMapping[int, Rational]):
for k in frozenset(k for k, v in m.items() if v == 0):
del m[k]


def _pf_rec(node: ast.expr):
match node:
case ast.BinOp(left=n_left, op=ast.Pow(), right=n_right):
res_base = _processSubnode(n_left)
exp = _processExponent(n_right)
_mul(res_base, exp)
return res_base
case ast.BinOp(left=n_left, op=op, right=n_right):
res_left = _processSubnode(n_left)
res_right = _processSubnode(n_right)
match op:
case ast.Mult():
Counter.update(res_left, res_right)
return res_left
case ast.Div():
Counter.subtract(res_left, res_right)
return res_left
case _:
raise ValueError(f"Invalid operator: {op.__class__.__name__}")
case _:
raise ValueError(f"Invalid expression: {ast.unparse(node)}")


def _processSubnode(node: ast.stmt) -> MutableMapping[int, Rational]:
match node:
case ast.Constant(value=val):
if not isinstance(val, Integral):
raise ValueError(f"{ast.unparse(node)} is not an integer.")

from primefac import primefac

res = {}
Counter.update(res, primefac(val))
return res
case _:
return _pf_rec(node)


def _processInteger(node: ast.expr) -> int:
neg = False
match node:
case ast.UnaryOp(op=ast.UAdd(), operand=node):
...
case ast.UnaryOp(op=ast.USub(), operand=node):
neg = True
match node:
case ast.Constant(value=val):
if not isinstance(val, Integral):
raise ValueError(f"{ast.unparse(node)} is not an integer.")
if neg:
val = -val
return val
case _:
raise ValueError(f"Invalid integer value: {ast.unparse(node)}")


def _processExponent(node: ast.expr) -> int | Q:
neg = False
match node:
case ast.UnaryOp(op=ast.UAdd(), operand=node):
...
case ast.UnaryOp(op=ast.USub(), operand=node):
neg = True
match node:
case ast.Constant(value=exp):
# integer exponent
if neg:
exp = -exp
return exp
case ast.BinOp(left=n_left, op=ast.Div(), right=n_right):
# fractional exponent
p = _processInteger(n_left)
q = _processInteger(n_right)
if neg:
p = -p
return Q(p, q)
case _:
raise ValueError(f"Invalid exponent: {ast.unparse(node)}")


def pfMapping(src: str):
node = ast.parse(src, mode="eval").body
res = _pf_rec(node)
_cleanZeroes(res)
return res


if __name__ == "__main__":
from viztracer import VizTracer

with VizTracer():
expression = "(3/2) * (3**4 / (20 * 4))**(-1/4)"
# desired result: {5: 1/4}
print(pfMapping(expression))
print(pfMapping(expression))

Continue reading...