Factoring polys in sympy

Question

I'm doing a very simple probability calculations of getting subset of X, Y, Z from set of A-Z (with corresponding probabilities x, y, z).

And because of very heavy formulas, in order to handle them, I'm trying to simplify (or collect or factor - I dont know the exact definition) these polynomial expressions using sympy.

So.. having this (a very simple probability calculation expression of getting subset of X,Y,Z from set of A-Z with corresponding probabilities x, y, z)

import sympy as sp

x, y, z = sp.symbols('x y z')

expression = (
    x * (1 - x) * y * (1 - x - y) * z +
    x * (1 - x) * z * (1 - x - z) * y +

    y * (1 - y) * x * (1 - y - x) * z +
    y * (1 - y) * z * (1 - y - z) * x +

    z * (1 - z) * y * (1 - z - y) * x +
    z * (1 - z) * x * (1 - z - x) * y
)

I want to get something like this

x * y * z * (6 * (1 - x - y - z) + (x + y) ** 2 + (y + z) ** 2 + (x + z) ** 2)

a poly, rewritten in way to have as few operations (+, -, *, **, ...) as possible

---

I tried factor(), collect(), simplify(). But result differs from my expectations. Mostly I get

2*x*y*z*(x**2 + x*y + x*z - 3*x + y**2 + y*z - 3*y + z**2 - 3*z + 3)

I know that sympy can combine polynomials into simple forms:

sp.factor(x**2 + 2*x*y + y**2)  # gives (x + y)**2

But how to make sympy to combine polynomials from expressions above?

---

If this is impossible task in sympy, may be there are any other options?

Answer 1

Putting together some of the methods happens to give a nice answer this time. It would be interesting to see if this strategy works more often than not on the equations you generate or if, as the name implies, this is just a lucky result this time.

def iflfactor(eq):
    """Return the "I'm feeling lucky" factored form of eq."""
    e = Mul(*[horner(e) if e.is_Add else e for e in
        Mul.make_args(factor_terms(expand(eq)))])
    r, e = cse(e)
    s = [ri[0] for ri in r]
    e = Mul(*[collect(ei.expand(), s) if ei.is_Add else ei for ei in
        Mul.make_args(e[0])]).subs(r)
    return e

>>> iflfactor(eq)  # using your equation as eq
2*x*y*z*(x**2 + x*y + y**2 + (z - 3)*(x + y + z) + 3)
>>> _.count_ops()
15

BTW, a difference between factor_terms and gcd_terms is that factor_terms will work harder to pull out common terms while retaining the original structure of the expression, very much like you would do by hand (i.e. looking for common terms in Adds that can be pulled out).

>>> factor_terms(x/(z+z*y)+x/z)
x*(1 + 1/(y + 1))/z
>>> gcd_terms(x/(z+z*y)+x/z)
x*(y*z + 2*z)/(z*(y*z + z))

For what it's worth,

Chris