ベアストウ法を実装した

以下の高次方程式の解を全て求められる。

試しに

を解いてみると

>>> solve([16, -152, 324, 162, 80, 50])
[(1.7058483573106117e-15+0.5000000000000019j), (1.7058483573106117e-15-0.5000000000000019j), 4.999999960549872, -0.5000000000000036, 5.000000039450128]

となり、複素数解や重解が求められているのがわかる。また、第二引数に精度が指定できるようになっている(デフォルトでは1e-9)。

>>> solve([16, -152, 324, 162, 80, 50], 1e-15)
[(-1.5133374263982614e-17+0.5j), (-1.5133374263982614e-17-0.5j), 4.999999985745681, -0.5, 5.000000014254319]

# -*- coding: utf-8 -*-

import math

def cutzero(a):
    c = 0
    while a[c] == 0:
        c += 1
    return a[c:]

def solve(a, eps=1e-9):
    a = cutzero(a)
    n = len(a) - 1
    solutions = []
    p0 = q0 = 0.0
    while n > 2:
        b = [0 for i in range(n+1)]
        c = [0 for i in range(n+1)]
        d = [0 for i in range(n+1)]
        b[0], b[1] = a[0], a[1]-p0*a[0]
        for i in range(2, n+1):
            b[i] = a[i] - p0*b[i-1] - q0*b[i-2]
        c[0], c[1] = 0, -b[0]
        for i in range(2, n+1):
            c[i] = -b[i-1] - p0*c[i-1] - q0*c[i-2]
        for i in range(2, n+1):
            d[i] = -b[i-2] - p0*d[i-1] - q0*d[i-2]
        D = c[n-1]*d[n] - d[n-1]*(c[n]+b[n-1])
        dp = (b[n]*d[n-1] - b[n-1]*d[n]) / D
        dq = (b[n-1]*(c[n]+b[n-1]) - b[n]*c[n-1]) / D
        if abs(dp) < eps and abs(dq) < eps:
            rt = p0*p0 - 4*q0
            imag = 1
            if rt < 0: imag = 1j
            rt = abs(rt)
            sa = (-p0 + math.sqrt(rt) * imag) / 2
            sb = (-p0 - math.sqrt(rt) * imag) / 2
            solutions.append(sa)
            solutions.append(sb)
            a = b[:n-1]
            n -= 2
        else:
            p0 += dp
            q0 += dq
    if n == 2:
        rt = a[1]*a[1] - 4*a[0]*a[2]
        imag = 1
        if rt < 0: imag = 1j
        rt = abs(rt)
        sa = (-a[1] + math.sqrt(rt) * imag) / (2*a[0])
        sb = (-a[1] - math.sqrt(rt) * imag) / (2*a[0])
        solutions.append(sa)
        solutions.append(sb)
    if n == 1:
        solutions.append(-a[1]/a[0])
    return solutions