Every Number Whose Power Is z

complex numbers · polar form · De Moivre's formula · n-th roots · ordering

In polar form, z = r\,e^{i\theta} with length r = |z| and angle \theta. Powers are easy in this form (De Moivre's formula):

\left(R\,e^{i\alpha}\right)^n = R^n e^{in\alpha}

the length is raised to the power n and the angle is multiplied by n. Run it backwards and you get the numbers w with w^n = z, the n-th roots of z: their length is R = r^{1/n}, and their angle \alpha must satisfy n\alpha = \theta up to a whole number of turns, so there are n different ones. Python's z ** (1 / n) returns just one of them.

Write all_roots(z, n) that returns the list of all n complex numbers w with w ** n == z (up to rounding), sorted by their angle measured anticlockwise from the positive real axis in the range [0, 2\pi), smallest angle first.

Examples

Input:  z = -8, n = 3
Output: [(1+1.7320508075688772j), (-2+0j), (1-1.7320508075688772j)]
Explanation: |-8| = 8 and 8 ** (1/3) = 2. The angles of the three cube roots are
pi/3, pi and 5pi/3; each one times 3 is pi up to whole turns, the angle of -8.

Input:  z = -1j, n = 2
Output: [(-0.7071067811865476+0.7071067811865476j), (0.7071067811865476-0.7071067811865476j)]
Explanation: the square roots of -1j have angles 3pi/4 and 7pi/4, in that order.

Values like 2.4e-16 where the exact answer is 0 count as correct.

Constraints

  • z != 0, and no z lies within 1e-6 (in angle) below the positive real axis
  • 1 <= n <= 1000
  • answers are compared with a tolerance of 1e-6; do not round

Goals

  • Use De Moivre's formula: raising to the power n multiplies the angle by n
  • Find all n complex n-th roots of a number from its polar form
  • List results in a stated order instead of trusting ** (1/n)
Starting Python…