rootsfinding

from math import*
from numbers import*
def roots(pol):
    pol=[float(n) for n in pol]
    deg=max([i for i in range(len(pol)) if pol[i]<>0])
    if deg==0: return None
    elif deg==1: return [-pol[0]/pol[1]]
    elif deg==2:
        a,b,c=pol[2],pol[1],pol[0]
        D=b*b-4*a*c
        if D<0: return [(-b-(0+1j)*(-D)**0.5)/(2*a), (-b+(0+1j)*(-D)**0.5)/(2*a)]
        else: return [(-b-D**0.5)/(2*a), (-b+D**0.5)/(2*a)]
    #Given ax^3+bx^2+cx+d=0. Reduce to x^3+ax^2+bx+c=0. Substitute x=y-a/3.
    #Reduce to y^3+yp+q=0. Substitute y=z-p/3z. Got z^6+qz^3-p^3/27=0
    elif deg==3: #kinda mystics why half of solutions can be omitted
        a,b,c=pol[2]/pol[3],pol[1]/pol[3],pol[0]/pol[3]
        p,q=b-a*a/3, c+2*a*a*a/27-a*b/3
        unity=[cos(2*pi*i/3)+sin(2*pi*i/3)*(0+1j) for i in range(3)]
        solve_z=(roots([-p*p*p/27, q, 1])[0]+0j)**(1./3) #solved for z
        zs=[solve_z*unity[i] for i in range(3)]
        ys=[z-p/(3*z) for z in zs]
        return [y-a/3 for y in ys]
    elif deg==4: #Reduce to y^4=py^2+qy+r
        a,b,c,d=pol[3]/pol[4],pol[2]/pol[4],pol[1]/pol[4], pol[0]/pol[4]
        p,q,r=3*a*a/8-b, a*b/2-a*a*a/8-c,3*a**4/256-a*a*b/16+a*c/4-d
        #print p,q,r, [4*p*r-q*q,8*r,4*p,8]
        D=roots([4*p*r-q*q,8*r,4*p,8]) #solving cubics for a constant :)
        if D[0].imag<D[1].imag: t=D[0].real
        else: t=D[1].real #choosing real root
        pre=((2*t+p+0j)**0.5).real
        return [x-a/4 for x in roots([t+pre*q/(2*(2*t+p)),pre,1])+roots([t-pre*q/(2*(2*t+p)),-pre,1])]
    else:
        print "I can't solve for big degrees cuz I don't know how to approximate complex roots efficiently"
        return None

print roots([1,-3,0,4,1])
#print roots([-1.0, 40.0, 24.0, 8])