-
- # Python 3. May also work in previous Python versions.
- # Inspired by https://bbs.archlinux.org/viewtopic.php?id=165527
-
- def interpret(tree):
- '''
- Expand a tree containing a recursively defined mathematical expression.
-
- Essentially the problem is turning the following expression into a sum of
- simple terms:
-
- -u - v*w
-
- where u, v and w may each be subexpressions of the same form, or simple
- variables.
-
- This function expects such an expression as a 3-element list, of which
- each element may be another list (recursively, hence the tree) or a
- 2-tuple. It returns an iterator which yields each term of the expanded
- expression one by one. A "term" here is a dictionary with two entries:
- "coef", the numeric coefficient, and "vars", the list of variables
- multiplied together to give the term.
-
- Like terms are not combined; in other words, if the input contains a
- subexpression like (A + AB)(C + BC), the result will contain (AC + ABC +
- ABC + ABBC), not (AC + 2ABC + ABBC). Multiplication is not treated as
- commutative, so variables are not rearranged within a term.
-
- Example usage:
-
- >>> expr = [(1, 0), (2, 1), [(0, 1), (1, 2), (2, 0)]]
- >>> for term in interpret(expr):
- ... print('{coef}*{vars}'.format(**term))
- ...
- -1*[(1, 0)] # -1*a_1s
- 1*[(2, 1), (0, 1)] # 1*a_21*a_s1
- 1*[(2, 1), (1, 2), (2, 0)] # 1*a_21*a_12*a_2s
- '''
- # The basic notion here is that every call to interpret() yields an
- # iterator over a sequence of terms. Therefore, recursive calls to
- # interpret() can be folded into the "result so far" by multiplying each
- # term in the first subexpression by -1, and multiplying each term in the
- # second subexpression by each term in the third subexpression (i.e.
- # applying the distributive property) and the result by -1.
-
- if len(tree) == 2:
- # Base case is easy: just yield one term containing the only variable
- # with a coefficient of 1.
- var = (tree[0], tree[1])
- yield dict(coef=1, vars=[var])
- elif len(tree) == 3:
- # tree[0] is our first subexpression. interpret(tree[0]) is a bunch of
- # terms (implicitly, a sum). Multiplying this subexpression by -1 is
- # the same as multiplying each term individually by -1 and summing
- # them together afterwards.
- for t in interpret(tree[0]):
- yield dict(coef=-1*t['coef'], vars=t['vars'])
- # tree[1] and tree[2] have to be multiplied together to give simple
- # terms that we can add to the result. Since interpret(tree[1]) and
- # interpret(tree[2]) are both iterable, we can do this in a nested
- # loop.
- for s in interpret(tree[1]):
- for t in interpret(tree[2]):
- # Two terms multiplied together is the same as one term with
- # all of their combined variables (in order) and the product
- # of their coefficients
- new_vars = s['vars'] + t['vars']
- new_coef = -1*s['coef']*t['coef']
- yield dict(coef=new_coef, vars=new_vars)
- else:
- raise ValueError("interpret(): malformed argument list")
-