hensel_lifting.py

from typing import List, Optional


def calc_f(f: List[int], x: int, m: Optional[int]=None) -> int:
    ret = 0
    for i, a in enumerate(f):
        term = a * x**i
        if m is not None:
            term %= m

        ret += term

    if m is not None:
        ret %= m

    return ret


def diff_f(f: List[int]) -> List[int]:
    ret = []
    for i, a in enumerate(f):
        if i == 0:
            continue
        ret.append(i*a)

    return ret


p = 521

# prepare polynomial and root (from sage)
f = [150, 323, 125, 449, 363, 43, 0, 106, 6, 391, 31]
x = 80

assert calc_f(f, x, p) == 0

# df/dx
_f = diff_f(f)
print(_f)
assert calc_f(_f, x, p) != 0

# lift to Z/p^2Z
b = calc_f(f, x) // p
d = calc_f(_f, x) % p
m = -b * pow(d, -1, p) % p
y = x + m*p

assert calc_f(f, y, p*p) == 0

# more lifting!!
# p^k -> p^{k+1}
for k in range(1, 100):
    pk = p**k
    pk1 = p**(k+1)

    b = calc_f(f, x) // pk
    d = calc_f(_f, x) % pk
    m = -b * pow(d, -1, pk) % pk
    y = x + m*pk

    res = calc_f(f, y, pk1)
    assert res == 0

    x = y