Kannan enumeration in Python
Enumeration with preprocessing that solves SVP in n^O(n) time.
- File
kannan_enumeration.py, 183 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/kannan-enumeration/. The authors are listed in prov/NY4D49.jsonl.
from fractions import Fraction
from math import floor, isqrt
from typing import Optional, Sequence, Union
Number = Union[int, Fraction]
def _dot(u: Sequence[Fraction], v: Sequence[Fraction]) -> Fraction:
return sum((a * b for a, b in zip(u, v)), Fraction(0))
def _gso(b: list[list[Fraction]]) -> tuple[list[list[Fraction]], list[Fraction]]:
"""Gram-Schmidt coefficients mu[i][j] (j < i) and squared norms N_i; ValueError if dependent."""
n = len(b)
stars: list[list[Fraction]] = []
norms: list[Fraction] = []
mu = [[Fraction(0)] * n for _ in range(n)]
for i, row in enumerate(b):
w = row[:]
for j in range(i):
mu[i][j] = _dot(row, stars[j]) / norms[j]
w = [a - mu[i][j] * c for a, c in zip(w, stars[j])]
nw = _dot(w, w)
if nw == 0:
raise ValueError("the basis vectors are linearly dependent")
stars.append(w)
norms.append(nw)
return mu, norms
def _lll_from(b: list[list[Fraction]], lo: int, delta: Fraction = Fraction(3, 4)) -> None:
"""LLL in place on the projected rows pi_lo(b_lo), ..., pi_lo(b_{n-1}).
Rows before lo are not changed; rows from lo on are size-reduced against all earlier rows,
and swaps happen only among positions >= lo.
"""
n = len(b)
k = lo + 1
while k < n:
mu, norms = _gso(b[: k + 1])
for j in range(k - 1, -1, -1):
q = floor(mu[k][j] + Fraction(1, 2))
if q:
b[k] = [a - q * c for a, c in zip(b[k], b[j])]
for l in range(j):
mu[k][l] -= q * mu[j][l]
mu[k][j] -= q
if norms[k] + mu[k][k - 1] ** 2 * norms[k - 1] >= delta * norms[k - 1]:
k += 1
else:
b[k - 1], b[k] = b[k], b[k - 1]
k = max(k - 1, lo + 1)
def _block_svp(
mu: list[list[Fraction]], norms: list[Fraction], k: int, h: int, stats: Optional[dict] = None
) -> tuple[list[int], Fraction]:
"""Shortest nonzero vector of pi_k(L(b_k..b_h)) by exact enumeration (iterative).
Returns block coefficients x (length h - k + 1) and the squared length
sum_j (x_j + sum_{i>j} x_i mu_ij)^2 N_j. The radius shrinks to the best length found.
"""
d = h - k + 1
best_x = [1] + [0] * (d - 1)
best = norms[k]
x = [0] * d
c = [Fraction(0)] * d
rho = [Fraction(0)] * (d + 1)
i, new = d - 1, True
while i < d:
if new:
c[i] = -sum((x[j] * mu[k + j][k + i] for j in range(i + 1, d)), Fraction(0))
q = max(best - rho[i + 1], Fraction(0)) / norms[k + i]
x[i] = floor(c[i]) - isqrt(floor(q)) - 1
else:
x[i] += 1
while (x[i] - c[i]) ** 2 * norms[k + i] + rho[i + 1] > best and x[i] <= c[i]:
x[i] += 1
val = rho[i + 1] + (x[i] - c[i]) ** 2 * norms[k + i]
if val > best:
i, new = i + 1, False
continue
rho[i] = val
if stats is not None:
stats["nodes"] += 1
if i == 0:
if val < best and any(x):
best, best_x = val, x[:]
new = False
else:
i, new = i - 1, True
return best_x, best
def _ext_gcd(a: int, b: int) -> tuple[int, int, int]:
"""(g, s, t) with s a + t b = g = gcd(a, b) >= 0."""
s0, s1, t0, t1 = 1, 0, 0, 1
while b:
q = a // b
a, b = b, a - q * b
s0, s1 = s1, s0 - q * s1
t0, t1 = t1, t0 - q * t1
if a < 0:
a, s0, t0 = -a, -s0, -t0
return a, s0, t0
def _insert(b: list[list[Fraction]], k: int, x: list[int]) -> None:
"""Replace b_k..b_{k+d-1} by a basis of the same lattice whose first vector is sum x_i b_{k+i}.
Needs gcd(x) = 1. Each step applies a 2 x 2 unimodular matrix to adjacent rows.
"""
x = x[:]
for j in range(len(x) - 2, -1, -1):
if x[j] == 0 and x[j + 1] == 0:
continue
g, s, t = _ext_gcd(x[j], x[j + 1])
p, q = x[j] // g, x[j + 1] // g
u, w = b[k + j], b[k + j + 1]
b[k + j] = [p * a + q * c for a, c in zip(u, w)]
b[k + j + 1] = [-t * a + s * c for a, c in zip(u, w)]
x[j], x[j + 1] = g, 0
def _size_reduce(b: list[list[Fraction]]) -> None:
"""Make |mu_ij| <= 1/2 for all j < i without changing the Gram-Schmidt vectors."""
for i in range(1, len(b)):
mu, _ = _gso(b[: i + 1])
for j in range(i - 1, -1, -1):
q = floor(mu[i][j] + Fraction(1, 2))
if q:
b[i] = [a - q * c for a, c in zip(b[i], b[j])]
for l in range(j):
mu[i][l] -= q * mu[j][l]
mu[i][j] -= q
def kannan_enumeration(basis: Sequence[Sequence[Number]], stats: Optional[dict] = None) -> list[list[Fraction]]:
"""Kannan's algorithm: an HKZ-reduced basis of the same lattice, in exact arithmetic.
The first row of the result is a shortest nonzero lattice vector. The recursion of
Kannan's algorithm on the projected lattices pi_k(L) runs on an explicit task stack:
(k, start) LLL-reduces the projected rows k..n-1 and then HKZ-reduces pi_{k+1};
(k, check) makes the block quasi-HKZ (N_k <= 4 N_{k+1}) by swaps, then enumerates a
shortest vector of pi_k(L) with radius N_k and inserts it. If stats is a dict, it gets
the counts 'nodes' (enumeration nodes) and 'swaps'. Raises ValueError on invalid input.
"""
b = [[Fraction(c) for c in row] for row in basis]
n = len(b)
if n == 0 or any(len(row) != len(b[0]) for row in b):
raise ValueError("need at least one basis row, all of the same length")
_gso(b)
if stats is not None:
stats.setdefault("nodes", 0)
stats.setdefault("swaps", 0)
tasks: list[tuple[int, bool]] = [(0, True)] # @line L1
while tasks: # @line L2
k, start = tasks.pop() # @line L3
if k >= n - 1:
continue
if start:
_lll_from(b, k) # @line L4
tasks += [(k, False), (k + 1, True)]
continue
mu, norms = _gso(b)
q = floor(mu[k + 1][k] + Fraction(1, 2)) # @line L5
if q:
b[k + 1] = [a - q * c for a, c in zip(b[k + 1], b[k])]
if norms[k] > 4 * norms[k + 1]: # @line L6
b[k], b[k + 1] = b[k + 1], b[k]
tasks += [(k, False), (k + 1, True)]
if stats is not None:
stats["swaps"] += 1
continue
mu, norms = _gso(b)
x, length = _block_svp(mu, norms, k, n - 1, stats) # @line L7
if length < norms[k]: # @line L8
_insert(b, k, x)
tasks.append((k + 1, True))
_size_reduce(b) # @line L9
return b
test_kannan_enumeration.py 255 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/kannan-enumeration/. The authors are listed in prov/NY4D49.jsonl.
import random
from fractions import Fraction
from math import floor, isqrt
import pytest
from kannan_enumeration import kannan_enumeration
def dot(u, v):
return sum((Fraction(a) * b for a, b in zip(u, v)), Fraction(0))
def combo(x, b):
return [sum((xi * row[j] for xi, row in zip(x, b)), Fraction(0)) for j in range(len(b[0]))]
def sub(u, v):
return [a - c for a, c in zip(u, v)]
def inverse(a):
"""Exact inverse of a square rational matrix (test-local, row reduction)."""
n = len(a)
m = [[Fraction(x) for x in row] + [Fraction(int(i == j)) for j in range(n)] for i, row in enumerate(a)]
for c in range(n):
p = next(r for r in range(c, n) if m[r][c] != 0)
m[c], m[p] = m[p], m[c]
m[c] = [x / m[c][c] for x in m[c]]
for r in range(n):
if r != c and m[r][c] != 0:
f = m[r][c]
m[r] = [x - f * y for x, y in zip(m[r], m[c])]
return [row[n:] for row in m]
def dual(b):
"""Rows d_i with <d_i, b_j> = [i = j] in the span of b: D = (B B^T)^-1 B."""
gi = inverse([[dot(u, v) for v in b] for u in b])
return [combo(gi[i], b) for i in range(len(b))]
def gso_norms(b):
bs, norms = [], []
for row in b:
w = [Fraction(c) for c in row]
for s, ns in zip(bs, norms):
if ns != 0:
w = sub(w, [dot(row, s) / ns * c for c in s])
bs.append(w)
norms.append(dot(w, w))
return bs, norms
def lll(b):
"""Textbook exact LLL with delta = 3/4 (test-local), used only to shrink oracle boxes."""
b = [[Fraction(c) for c in row] for row in b]
k = 1
while k < len(b):
bs, ns = gso_norms(b[: k + 1])
for j in range(k - 1, -1, -1):
q = floor(dot(b[k], bs[j]) / ns[j] + Fraction(1, 2))
if q:
b[k] = sub(b[k], [q * c for c in b[j]])
mu = dot(b[k], bs[k - 1]) / ns[k - 1]
nk = ns[k]
if nk >= (Fraction(3, 4) - mu * mu) * ns[k - 1]:
k += 1
else:
b[k], b[k - 1] = b[k - 1], b[k]
k = max(k - 1, 1)
return b
def sqrt_up(q):
"""A rational upper bound for the square root of a nonnegative rational."""
k = 10**6
return Fraction(isqrt(q.numerator * k * k // q.denominator) + 1, k)
def random_basis(rng, n, m, lo=-4, hi=4):
while True:
b = [[rng.randint(lo, hi) for _ in range(m)] for _ in range(n)]
_, ns = gso_norms(b)
if all(x != 0 for x in ns):
return b
def svp_sq(b):
"""Exact lambda_1^2 by pruned enumeration on a test-local LLL basis.
With Gram-Schmidt data of C, ||x C||^2 = sum_j (x_j + sum_{i>j} x_i mu_ij)^2 N_j, so the
coefficients can be fixed from the last to the first, keeping partial sums <= best.
"""
c = lll(b)
bs, ns = gso_norms(c)
n = len(c)
mu = [[dot(c[i], bs[j]) / ns[j] for j in range(n)] for i in range(n)]
best = [dot(c[0], c[0])]
def search(j, x, partial):
if j < 0:
if any(x):
best[0] = min(best[0], partial)
return
centre = -sum((x[i] * mu[i][j] for i in range(j + 1, n)), Fraction(0))
w = sqrt_up((best[0] - partial) / ns[j])
for xj in range(floor(centre - w), floor(centre + w) + 1):
add = (centre - xj) ** 2 * ns[j]
if partial + add <= best[0]:
x[j] = xj
search(j - 1, x, partial + add)
x[j] = 0
search(n - 1, [0] * n, Fraction(0))
return best[0]
def same_lattice(b, c):
"""C = U B with U integral and det U = +-1."""
g = inverse([[dot(u, v) for v in b] for u in b])
u = [[sum((dot(ci, b[k]) * g[k][j] for k in range(len(b))), Fraction(0)) for j in range(len(b))] for ci in c]
if any(x.denominator != 1 for row in u for x in row):
return False
return abs(det(u)) == 1
def det(a):
a = [row[:] for row in a]
n = len(a)
s = Fraction(1)
for c in range(n):
p = next((r for r in range(c, n) if a[r][c] != 0), None)
if p is None:
return Fraction(0)
if p != c:
a[c], a[p] = a[p], a[c]
s = -s
s *= a[c][c]
for r in range(c + 1, n):
f = a[r][c] / a[c][c]
a[r] = [x - f * y for x, y in zip(a[r], a[c])]
return s
def projected_block(c, k, h):
"""The vectors pi_k(c_i) for k <= i <= h (0-based, inclusive)."""
bs, ns = gso_norms(c)
out = []
for i in range(k, h + 1):
v = [Fraction(a) for a in c[i]]
for j in range(k):
f = dot(c[i], bs[j]) / ns[j]
v = [a - f * s for a, s in zip(v, bs[j])]
out.append(v)
return out
def is_size_reduced(c):
bs, ns = gso_norms(c)
return all(abs(dot(c[k], bs[j]) / ns[j]) <= Fraction(1, 2) for k in range(len(c)) for j in range(k))
def is_hkz(c):
_, ns = gso_norms(c)
n = len(c)
return is_size_reduced(c) and all(ns[k] == svp_sq(projected_block(c, k, n - 1)) for k in range(n))
def test_worked_example():
b = [[4, 0, 0], [-2, 3, 0], [1, 1, 3]]
stats = {}
c = kannan_enumeration(b, stats)
assert c == [[-1, -1, -3], [2, 3, 0], [-2, 3, 0]]
assert gso_norms(c)[1] == [11, Fraction(118, 11), Fraction(648, 59)]
assert stats["swaps"] == 0
assert same_lattice(b, c) and is_hkz(c)
def test_quasi_hkz_swaps():
b = [[80, 0, 0, 0, 0], [40, 58, 0, 0, 0], [0, 29, 42, 0, 0], [0, 0, 21, 30, 0], [0, 0, 0, 15, 22]]
stats = {}
c = kannan_enumeration(b, stats)
assert stats["swaps"] == 2
assert c[0] == [0, 0, 0, 15, 22] and dot(c[0], c[0]) == 709 == svp_sq(b)
assert same_lattice(b, c) and is_hkz(c)
def test_random_bases_are_hkz_reduced():
rng = random.Random(61)
count = 0
for n, reps in ((2, 80), (3, 60), (4, 35), (5, 10), (6, 4)):
for _ in range(reps):
m = n + rng.randint(0, 1)
b = random_basis(rng, n, m, -6, 6)
c = kannan_enumeration(b)
assert same_lattice(b, c)
assert is_hkz(c) # includes ||c_1||^2 = lambda_1^2
count += 1
assert count == 189
def test_skewed_bases():
rng = random.Random(62)
for _ in range(25):
n = rng.randint(3, 5)
diag = [rng.randint(30, 60)]
for _ in range(n - 1):
diag.append(max(2, diag[-1] * 3 // 4))
b = [[0] * n for _ in range(n)]
for i in range(n):
b[i][i] = 2 * diag[i]
if i:
b[i][i - 1] = diag[i - 1]
c = kannan_enumeration(b)
assert same_lattice(b, c) and is_hkz(c)
def test_already_reduced_and_idempotent():
eye = [[int(i == j) for j in range(5)] for i in range(5)]
assert kannan_enumeration(eye) == eye
rng = random.Random(63)
for _ in range(30):
n = rng.randint(1, 5)
c = kannan_enumeration(random_basis(rng, n, n, -5, 5))
assert gso_norms(kannan_enumeration(c))[1] == gso_norms(c)[1]
def test_dimension_one():
assert kannan_enumeration([[-7, 3]]) == [[-7, 3]]
assert kannan_enumeration([[Fraction(5, 2)]]) == [[Fraction(5, 2)]]
def test_nearly_dependent():
rng = random.Random(64)
for _ in range(20):
n = rng.randint(2, 4)
base = [rng.randint(1, 9) for _ in range(n)]
b = [[x * 1000 + (i == j) for j, x in enumerate(base)] for i in range(n)]
c = kannan_enumeration(b)
assert same_lattice(b, c) and is_hkz(c)
b = [[1, 0], [1, Fraction(1, 1000)]]
c = kannan_enumeration(b)
assert dot(c[0], c[0]) == Fraction(1, 10**6) and is_hkz(c)
def test_invalid_input():
with pytest.raises(ValueError):
kannan_enumeration([])
with pytest.raises(ValueError):
kannan_enumeration([[1, 2], [2, 4]])
with pytest.raises(ValueError):
kannan_enumeration([[1, 2], [3]])
What the next level needs
- the tests did not run in this build