Search code examples
pythonsympypolynomials

How to isolate the coefficients of the terms in a polynomial in sympy?

In Python, I do some SymPy calculations that yield an array full of polynomials such as:

a*(a*(a*(a + b) + b*(a + b)) + b*(a*(a + b) + b*(a + b))) + b*(a*(a*(a + b) + b*(a + b)) + b*(a*(a + b) + b*(a + b)))

Note that this example happens to simplify to (a+b)**4, but this won't always be the case obviously. So how do I convert this expression to the following form:

c_1*a**4 + c_2*a**3*b + ... + c_n*b**4

And once I have such an expression, how would I extract the exponents c_1, ..., c_n? All I have is the .exp command, but it only works on expressions of the form a**n (i.e. no mixture of a and b and a coefficient of 1).

Any help would be majorly appreciated.

Solution

  • The Poly class is useful (running with isympy)

    from sympy import Poly
from sympy.abc import a, b

expr = a*(a*(a*(a + b) + b*(a + b)) + b*(a*(a + b) + b*(a + b))) + b*(a*(a*(a + b) + b*(a + b)) + b*(a*(a + b) + b*(a + b)))

poly = Poly(expr)
poly.as_expr()

    Output:

     4      3        2  2        3    4
a  + 4⋅a ⋅b + 6⋅a ⋅b  + 4⋅a⋅b  + b

    And do extract the coefficients c_1, ..., c_n, I would do:

    poly.coeffs()

    Output:

    [1, 4, 6, 4, 1]