python - 在 Python 中计算单位的 n 次根

Question

所以，我正在尝试编写一个算法 croot(k, n)，它返回 n == n 的第 k 个单位根。我得到的答案大多是正确的，但是它给了我非常奇怪的表示，对于某些数字来说似乎是错误的。这是一个例子。

import cmath

def croot(k, n):
    if n<=0:
        return None
    return cmath.exp((2 * cmath.pi * 1j * k) / n)


for k in range(8):
    print croot(k, 8)

输出是:

(1+0j)
(0.70710...+0.70710...j)
(6.12323399574e-17+1j)

哇哇哇。所以 k = 2 和 n = 8 时的根是错误的，因为它应该只是 i，可以表示为 1j、j 或 1.00000j 等。有人可以帮我吗？我这样做是因为我正在尝试编写 FFT 算法。我对复数和 Python 不是很有经验，所以我很可能会犯一个简单的错误。

谢谢，

如果你们需要任何其他信息，请询问。

Answer 1

下面是单位立方根和四次根的用法示例。输入数组应解释为多项式系数。

>>> import numpy as np
>>> np.roots([1, 0, 0, -1])
array([-0.5+0.8660254j, -0.5-0.8660254j,  1.0+0.j       ])
>>> np.roots([1, 0, 0, 0, -1])
array([ -1.00000000e+00+0.j,   5.55111512e-17+1.j,   5.55111512e-17-1.j,
         1.00000000e+00+0.j])

---

编辑: 多项式系数在输入数组 p 到 np.roots(p) 中给出按以下顺序:

p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n]

例如，要返回 n 单位根，即方程 1 * x**n - 1 == 0 的解，您将使用像 p = [1] + [0] * (n - 1) + [-1] 这样的输入。