Bostan-Mori algorithm in Python
Computes the n-th coefficient of a rational function by halving n and multiplying by the conjugate denominator.
- File
bostan_mori_algorithm.py, 188 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bostan-mori-algorithm/. The authors are listed in prov/8WEKFW.jsonl.
"""The Bostan-Mori algorithm: the N-th coefficient of a rational power series.
A polynomial a_0 + a_1 x + ... + a_n x^n is the list [a_0, ..., a_n], lowest
degree first; the zero polynomial is []. Coefficients may be int, Fraction, or
any ring element type with + - * and unary minus (and `/` for the inverse of a
unit). The constant term of the denominator must be a unit.
One step multiplies the numerator and the denominator by Q(-x). The new
denominator Q(x) Q(-x) has only even powers, so it is W(x^2), and the
coefficient of x^N comes only from the even or only from the odd part of the
new numerator. N is halved, so there are about log2 N steps, each with two
polynomial products. All loops are iterative.
"""
from __future__ import annotations
from typing import Any, Callable, List, Sequence
Poly = List[Any]
MulFn = Callable[[Sequence[Any], Sequence[Any]], Poly]
KARATSUBA_CUTOFF = 32
def schoolbook_mul(a: Sequence[Any], b: Sequence[Any]) -> Poly:
"""Return the full product of two coefficient lists, len(a) + len(b) - 1 entries."""
if not a or not b:
return []
c: Poly = [a[0] - a[0]] * (len(a) + len(b) - 1)
for i, x in enumerate(a):
for j, y in enumerate(b):
c[i + j] = c[i + j] + x * y
return c
def _add(a: Sequence[Any], b: Sequence[Any]) -> Poly:
"""Return the coefficientwise sum of two lists of possibly different lengths."""
if len(a) < len(b):
a, b = b, a
return [x + y for x, y in zip(a, b)] + list(a[len(b):])
def karatsuba_mul(a: Sequence[Any], b: Sequence[Any]) -> Poly:
"""Return the full product of two coefficient lists by Karatsuba's method.
Works over every commutative ring. The recursion splits the longer operand
in halves, so its depth is O(log(len(a) + len(b))).
"""
if not a or not b:
return []
if len(a) < len(b):
a, b = b, a
if len(b) <= KARATSUBA_CUTOFF:
return schoolbook_mul(a, b)
zero = a[0] - a[0]
h = len(a) // 2
a0, a1 = list(a[:h]), list(a[h:])
if len(b) <= h:
# Unbalanced: b fits in the low half; multiply each half of a by b.
lo = karatsuba_mul(a0, b)
hi = karatsuba_mul(a1, b)
c = [zero] * (len(a) + len(b) - 1)
for i, v in enumerate(lo):
c[i] = c[i] + v
for i, v in enumerate(hi):
c[i + h] = c[i + h] + v
return c
b0, b1 = list(b[:h]), list(b[h:])
z0 = karatsuba_mul(a0, b0)
z2 = karatsuba_mul(a1, b1)
sa = _add(a0, a1)
sb = _add(b0, b1)
z1 = karatsuba_mul(sa, sb)
for i, v in enumerate(z0):
z1[i] = z1[i] - v
for i, v in enumerate(z2):
z1[i] = z1[i] - v
c = [zero] * (len(a) + len(b) - 1)
for i, v in enumerate(z0):
c[i] = c[i] + v
for i, v in enumerate(z1):
if i + h < len(c):
c[i + h] = c[i + h] + v
for i, v in enumerate(z2):
c[i + 2 * h] = c[i + 2 * h] + v
return c
def _unit_inverse(c: Any) -> Any:
"""Return c^(-1), or raise ValueError if c is not a unit."""
if isinstance(c, bool):
raise TypeError("bool is not a ring element here")
if isinstance(c, int):
if c in (1, -1):
return c
raise ValueError("not a unit of Z")
if c == 0:
raise ValueError("zero is not a unit")
one = c - c + 1
try:
return one / c
except ZeroDivisionError as exc:
raise ValueError("not a unit") from exc
def normalize(a: Sequence[Any]) -> Poly:
"""Return a copy of a without zero entries at the top."""
out = list(a)
while out and out[-1] == 0:
out.pop()
return out
def bostan_mori(p: Sequence[Any], q: Sequence[Any], n: int, mul: MulFn = karatsuba_mul, steps: list = None) -> Any:
"""Return the coefficient of x^n in the power series p / q.
Precondition: q(0) is a unit. Any degrees are allowed. Cost: two products
of polynomials with at most deg q + 1 coefficients per halving of n, once
deg p < deg q.
"""
p = normalize(p)
q = normalize(q)
if not q:
raise ZeroDivisionError("the denominator is zero")
zero = q[0] - q[0]
if n < 0 or not p:
return zero
_unit_inverse(q[0]) # fail early if q(0) is not a unit
while n >= 1: # @line L1
qm = [c if i % 2 == 0 else -c for i, c in enumerate(q)] # @line L2
u = mul(p, qm) # @line L3
v = mul(q, qm) # @line L4
p = u[n % 2 :: 2] # @line L5
q = v[0::2] # @line L6
n //= 2 # @line L7
p = normalize(p)
q = normalize(q)
if steps is not None:
steps.append((n, list(p), list(q)))
if not p:
return zero
c0 = p[0] if p else zero # @line L8
return c0 * _unit_inverse(q[0])
def nth_term(c: Sequence[Any], a: Sequence[Any], n: int, mul: MulFn = karatsuba_mul) -> Any:
"""Return a_n for a_i = c_1 a_(i-1) + ... + c_d a_(i-d), given a_0, ..., a_(d-1).
c = [c_1, ..., c_d] and a = [a_0, ..., a_(d-1)] must have the same length d.
"""
d = len(c)
if len(a) != d:
raise ValueError("need exactly d initial values")
if n < 0:
raise ValueError("n must be nonnegative")
if d == 0:
return 0
if n < d:
return a[n]
one = c[0] - c[0] + 1
q = [one] + [-x for x in c]
p = mul(list(a), q)[:d]
return bostan_mori(p, q, n, mul)
def series_coefficient_naive(p: Sequence[Any], q: Sequence[Any], n: int) -> Any:
"""Reference: the first n + 1 coefficients of p / q by long division of series."""
q = normalize(q)
inv = _unit_inverse(q[0])
zero = q[0] - q[0]
s: Poly = []
for i in range(n + 1):
t = p[i] if i < len(p) else zero
for j in range(1, min(i, len(q) - 1) + 1):
t = t - q[j] * s[i - j]
s.append(t * inv)
return s[n]
if __name__ == "__main__":
st = []
print(bostan_mori([0, 1], [1, -1, -1], 10, steps=st))
for row in st:
print(row)
print(nth_term([1, 1], [0, 1], 90))
test_bostan_mori_algorithm.py 194 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/bostan-mori-algorithm/. The authors are listed in prov/8WEKFW.jsonl.
import random
from fractions import Fraction
import pytest
from bostan_mori_algorithm import bostan_mori, nth_term, schoolbook_mul, series_coefficient_naive
class ModP:
"""An element of Z/pZ, enough for the tests."""
def __init__(self, v: int, p: int) -> None:
self.p = p
self.v = v % p
def _c(self, o):
return o.v if isinstance(o, ModP) else o
def __add__(self, o):
return ModP(self.v + self._c(o), self.p)
__radd__ = __add__
def __sub__(self, o):
return ModP(self.v - self._c(o), self.p)
def __rsub__(self, o):
return ModP(self._c(o) - self.v, self.p)
def __mul__(self, o):
return ModP(self.v * self._c(o), self.p)
__rmul__ = __mul__
def __neg__(self):
return ModP(-self.v, self.p)
def __truediv__(self, o):
d = self._c(o) % self.p
if d == 0:
raise ZeroDivisionError
return ModP(self.v * pow(d, -1, self.p), self.p)
def __rtruediv__(self, o):
return ModP(o, self.p) / self
def __eq__(self, o):
return self.v == self._c(o) % self.p
def __hash__(self):
return hash(self.v)
def __repr__(self):
return f"{self.v} mod {self.p}"
# Independent oracles: run the recurrence forward, or power the companion
# matrix by repeated squaring modulo p.
def run_recurrence(c, a, n):
seq = list(a)
d = len(c)
while len(seq) <= n:
seq.append(sum(c[i] * seq[-1 - i] for i in range(d)))
return seq[n]
def companion_power_term(c, a, n, p):
d = len(c)
if n < d:
return a[n] % p
# state (a_i, ..., a_(i+d-1)) -> (a_(i+1), ..., a_(i+d))
m = [[0] * d for _ in range(d)]
for i in range(d - 1):
m[i][i + 1] = 1
for j in range(d):
m[d - 1][j] = c[d - 1 - j] % p
def mm(x, y):
return [[sum(x[i][t] * y[t][j] for t in range(d)) % p for j in range(d)] for i in range(d)]
r = [[int(i == j) for j in range(d)] for i in range(d)]
e = n
b = m
while e:
if e & 1:
r = mm(r, b)
b = mm(b, b)
e >>= 1
return sum(r[0][j] * a[j] for j in range(d)) % p
PRIMES = [2, 3, 7, 998244353]
def test_page_example():
steps = []
assert bostan_mori([0, 1], [1, -1, -1], 10, steps=steps) == 55
assert steps == [
(5, [0, 1], [1, -3, 1]),
(2, [1, 1], [1, -7, 1]),
(1, [1, 8], [1, -47, 1]),
(0, [55, 8], [1, -2207, 1]),
]
assert nth_term([1, 1], [0, 1], 90) == 2880067194370816120
def test_edge_cases():
assert bostan_mori([], [1, 1], 5) == 0
assert bostan_mori([3], [1], 0) == 3
assert bostan_mori([3], [1], 4) == 0
assert bostan_mori([1, 2, 3], [1], 2) == 3
assert bostan_mori([1, 2, 3], [1, 0, 0], 2) == 3
assert bostan_mori([1, 2, 3, 0, 0], [1, -1, 0, 0], 7) == 6
assert bostan_mori([1], [1, -1], 0) == 1
assert bostan_mori([1], [1, -1], 10**30) == 1
assert bostan_mori([1], [-1, 1], 3) == -1
assert bostan_mori([5, 7], [1, -2], -1) == 0
# numerator of higher degree than the denominator
p, q = [1, 2, 3, 4, 5, 6, 7], [1, -1, 1]
for n in range(20):
assert bostan_mori(p, q, n) == series_coefficient_naive(p, q, n)
with pytest.raises(ValueError):
bostan_mori([1], [2, -1], 3)
with pytest.raises(ZeroDivisionError):
bostan_mori([1], [0, 0], 3)
assert nth_term([], [], 5) == 0
assert nth_term([2], [3], 0) == 3
assert nth_term([2], [3], 10) == 3 * 2**10
with pytest.raises(ValueError):
nth_term([1, 1], [0], 3)
def test_over_z_random():
rng = random.Random(1)
for _ in range(300):
d = rng.randint(1, 8)
c = [rng.randint(-3, 3) for _ in range(d)]
a = [rng.randint(-5, 5) for _ in range(d)]
n = rng.randint(0, 300)
assert nth_term(c, a, n) == run_recurrence(c, a, n)
p = [rng.randint(-5, 5) for _ in range(rng.randint(0, d + 3))]
q = [rng.choice([1, -1])] + [rng.randint(-3, 3) for _ in range(rng.randint(0, d))]
m = rng.randint(0, 120)
assert bostan_mori(p, q, m) == series_coefficient_naive(p, q, m)
@pytest.mark.parametrize("p", PRIMES)
def test_over_z_mod_p(p):
rng = random.Random(p)
for _ in range(60):
d = rng.randint(1, 10)
c = [rng.randrange(p) for _ in range(d)]
a = [rng.randrange(p) for _ in range(d)]
n = rng.randint(0, 400)
got = nth_term([ModP(x, p) for x in c], [ModP(x, p) for x in a], n)
assert got == run_recurrence(c, a, n) % p
for _ in range(20):
d = rng.randint(1, 6)
c = [rng.randrange(p) for _ in range(d)]
a = [rng.randrange(p) for _ in range(d)]
n = rng.randrange(10**18)
got = nth_term([ModP(x, p) for x in c], [ModP(x, p) for x in a], n)
assert got == companion_power_term(c, a, n, p)
# non-monic constant term, a unit mod p
for _ in range(20):
q = [ModP(rng.randrange(1, p), p)] + [ModP(rng.randrange(p), p) for _ in range(rng.randint(0, 5))]
num = [ModP(rng.randrange(p), p) for _ in range(rng.randint(0, 6))]
m = rng.randint(0, 60)
assert bostan_mori(num, q, m) == series_coefficient_naive(num, q, m)
def test_over_fractions():
rng = random.Random(5)
for _ in range(60):
q = [Fraction(rng.randint(1, 5), rng.randint(1, 4))] + [Fraction(rng.randint(-4, 4), rng.randint(1, 3)) for _ in range(rng.randint(0, 4))]
p = [Fraction(rng.randint(-4, 4), rng.randint(1, 3)) for _ in range(rng.randint(0, 5))]
n = rng.randint(0, 40)
assert bostan_mori(p, q, n) == series_coefficient_naive(p, q, n)
# 1 / (2 - x) = sum x^n / 2^(n + 1)
assert bostan_mori([Fraction(1)], [Fraction(2), Fraction(-1)], 20) == Fraction(1, 2**21)
def test_huge_index_over_z():
# Fibonacci identity F(2k) = F(k) (2 F(k + 1) - F(k)), checked at k = 5000
k = 5000
fk = nth_term([1, 1], [0, 1], k)
fk1 = nth_term([1, 1], [0, 1], k + 1)
assert nth_term([1, 1], [0, 1], 2 * k) == fk * (2 * fk1 - fk)
assert fk == run_recurrence([1, 1], [0, 1], k)
What the next level needs
- the tests did not run in this build