Search code examples
python-3.xsympypolynomials

Expanding polynomial with variable exponent in sympy

I am not very experienced with sympy so sorry if this is a simple question.

How can I use sympy to expand binomial expressions? For example say I want to have sympy compute the coefficient of $x^2$ in the polynomial $(x^2 + x + 1)^n$ (where I would expect the answer to be $n + \binom{n}{2}$).

I tried the following code:

x = symbols('x')
n = symbols('n', integer=True, nonnegative = True)
expand((x**2+x+1)**n)

but the result is just $(x^2+x+1)^n$ whereas I would want the binomial expansion, i.e $(1+x+x^2)^n=\sum_{i=0}^{2n}{\left(\sum_{l=\max{0,i-n}}^{[i/2]}{\binom{n}{i-2l,l,n-i+l}}\right)x^i}$.

Thanks in advance.

Solution

  • If the exponent is not symbolic then the following gives the coefficient very quickly for the power of an arbitrary polynomial with integer coefficients, e.g.,

    >>> eq
x**3 + 3*x + 2
>>> (Poly(eq)**42).coeff_monomial(x**57)
2294988464559317378977138572972

    But there is currently no routine to indicate the coefficient if the exponent of the polynomial is symbolic. rsolve can also be used to express the closed form if a pattern can be seen in the coefficient, too:

    >>> print([((x**2+x+1)**i).expand().coeff(x**2) for i in range(8)])
[0, 1, 3, 6, 10, 15, 21, 28]
>>> from sympy.abc import n
>>> f=Function('f') # f(n) represents the coefficient of x**2 for a given n

    The coefficient for x^2 for a given n is n more than the last value:

    >>> rsolve(f(n)-f(n-1)-n, f(n),{f(0):0,f(1):1})
n*(n + 1)/2

    This final expression is the coefficient of x^2 for arbitrary n.

    Issue 17889 gives a routine that will compute the coefficient of a term in a univariate polynomial (with arbitrary coefficients for each term) raised to the power of n:

    >>> eq = 2 + x + x**2
>>> unicoeff(eq, 4).simplify()
Piecewise(
    (0, n < 2),
    (2**(n - 3)*n*(n - 1), n < 3),
    (2**(n - 4)*n**2*(n - 1), n < 4),
    (2**n*n*(n - 1)*(n**2 + 19*n + 6)/384, True))
>>> _.subs(n, 5)
210
>>> (eq**5).expand().coeff(x**4)
210

    For your expression (where the constant is 1):

    >>> unicoeff(1+x+x**2,2).simplify()
Piecewise((0, n < 1), (n, n < 2), (n*(n + 1)/2, True))