Python – Collecting like term of an expression in Sympy

mathpolynomialspythonpython-2.7sympy

I am currently dealing with functions of more than one variable and need to collect like terms in an attempt to simplify an expression.

Say the expression is written as follows:

x = sympy.Symbol('x')
y = sympy.Symbol('y')
k = sympy.Symbol('k')
a = sympy.Symbol('a')

z = k*(y**2*(a + x) + (a + x)**3/3) - k((2*k*y*(a + x)*(n - 1)*(-k*(y**2*(-a + x) + (-a + x)**3/3) + k*(y**2*(a + x) + (a + x)**3/3)) + y)**2*(-a + k*(n - 1)*(y**2 + (a + x)**2)*(-k*(y**2*(-a + x)))))
zEx = z.expand()
print type(z)
print type(zEx)

EDIT: Formatting to add clarity and changed the expression z to make the problem easier to understand.

Say z contains so many terms, that sifting through them by eye. and selecting the appropriate terms, would take an unsatisfactory amount of time.

I want to collect all of the terms which are ONLY a multiple of a**1. I do not care for quadratic or higher powers of a, and I do not care for terms which do not contain a.

The type of z and zEx return the following:

print type(z)
print type(zEx)
>>>
<class 'sympy.core.add.Add'>
<class 'sympy.core.mul.Mul'>

Does anyone know how I can collect the terms which are a multiple of a , not a^0 or a^2?

tl'dr

Where z(x,y) with constants a and k described by z and zEx and their type(): How can one remove all non-a terms from z AND remove all quadratic or higher terms of a from the expression? Such that what is left is only the terms which contain a unity power of a.

Best Answer

In addition to the other answers given, you can also use collect as a dictionary.

print(collect(zEx,a,evaluate=False)[a])

yields the expression

k*x**2 + k*y**2
Related Topic