Need to create function in Python that finds all complex and real roots of z^n=x(so the nth-roots of x)

Question

beginner programmer here. I have been assigned to create a function 'Roots' that takes two parameters x and n(n has to be an integer) and then calculates all complex and real roots of the equation z^n=x. However, the only module/package I can use is math. Also, I have been told that the certain aspects of the following function 'Power_complex' play a big role into creating 'Roots':

def Power_complex(re, im, n):     #calculates the n-th power of a complex number(lets call this a), where 're' is the real part and 'im' the imaginary part
import math
r=math.sqrt((re)**2+(im)**2)      #calculates modulus
h=math.atan2(re,im)               #calculates argument(angle)
ren=(r**n)*math.cos(h*n)          #calculates the real part of a^n
imn=(r**n)*math.sin(h*n)          #calculates the imaginary part of a^n
return ren, imn
print '(',re, '+', im, 'i',')','^',n,'=',ren,'+',imn,'i' #displays the result

Also, I need to somehow implement a for loop into 'Roots'. I have been pondering over this for hours, but alas I really can't figure it out and am hoping one of you can help me out here.

BTW my python version is 2.7.10

Answer 1

The expression for the solutions is ( explained here ):

when

In the case that z^n is real, equal to the x in your question, then r = |x| and the angle is 0 or pi for positive and negative values, respectively.

So you make the modulus and argument as you have done, then make every solution corresponding to a value of k

z = [r**(1./n) * exp(1j * (theta + 2*pi*k) / n ) for k in range(n)]

This line uses a Python technique called list comprehension. An eqvivalent way of doing it (that you may be more familiar to) could be:

z = []
for k in range(n):
    nthroot = r**(1./n) * exp( 1j * (theta + 2*pi*k) / n )
    z.append(nthroot)

Printing them out could be done in the same fashion, using a for-loop:

for i in range(len(z)):
    print "Root #%d = %g + i*%g" % (i, z[i].real, z[i].imag)

Note that the exp -function used must be from the module cmath ( math can't handle complex numbers). If you are not allowed to use cmath , then I suggest you rewrite the expression for the solutions to the form without modulus and argument.