hensel module

Pure-Python implementation of Hensel lifting for square roots modulo a prime power.

hensel(root: int, prime: int, exponent: int = 1) → int

Lift a square root of a value modulo prime ** exponent to the square root of that same value modulo prime ** (exponent + 1).

More specifically, let square be a nonnegative integer that is the least nonnegative residue of the congruence class root ** 2 modulo prime ** exponent. Use Hensel lifting to return an integer that represents the square root modulo prime ** (exponent + 1) of the congruence class represented by the integer square modulo prime ** (exponent + 1).

>>> hensel(4, 7)
39
>>> hensel(2, 7, 2)
2

This function implements a lifting operation even for those cases in which the root has the supplied prime as a factor (or is zero).

>>> hensel(28, 7, 3)
273
>>> pow(28, 2, 7 ** 3) == pow(273, 2, 7 ** 4)
True
>>> hensel(256, 2, 12)
512
>>> pow(256, 2, 2 ** 12) == pow(512, 2, 2 ** 13)
True

This function lifts distinct roots to distinct roots when possible.

>>> def roots(s, m):
...     return [r for r in range(0, m) if pow(r, 2, m) == s]
>>> [hensel(r, 3, 5) for r in roots(81, 3 ** 5)] == roots(81, 3 ** 6)
True

However, when the root has the supplied prime as a factor, it may be the case that not all roots modulo prime ** (exponent + 1) can be obtained via lifting. In that case, the number of distinct roots that can be obtained is equivalent to the number of distinct roots that are available to lift.

>>> [hensel(r, 2, 5) for r in roots(16, 2 ** 5)]
[12, 28, 44, 60]
>>> roots(16, 2 ** 6)
[4, 12, 20, 28, 36, 44, 52, 60]

Any attempt to invoke this function with arguments that do not have the expected types (or do not fall within the supported ranges) raises an exception. If prime is not a prime number, the behavior of this function is not specified.

>>> hensel('abc', 7)
Traceback (most recent call last):
  ...
TypeError: 'str' object cannot be interpreted as an integer
>>> hensel(2, {})
Traceback (most recent call last):
  ...
TypeError: 'dict' object cannot be interpreted as an integer
>>> hensel(2, 7, [])
Traceback (most recent call last):
  ...
TypeError: 'list' object cannot be interpreted as an integer
>>> hensel(2, -1)
Traceback (most recent call last):
  ...
ValueError: prime must be a positive integer
>>> hensel(2, 7, -1)
Traceback (most recent call last):
  ...
ValueError: exponent must be a nonnegative integer

The examples below verify the correct behavior of the function on a range of different inputs.

>>> all(
...     pow(r, 2, p ** k) == pow(hensel(r, p, k), 2, p ** (k + 1))
...     for k in range(0, 5)
...     for p in [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31]
...     for r in range(0, p ** k)
... )
True
>>> all(
...     lifted.issubset(actual) and (
...         len(actual) == len(lifted)
...         or
...         len(actual) == len(lifted) * p
...     )
...     for k in range(0, 4)
...     for p in [2, 3, 5, 7, 11, 13]
...     for s in [pow(x, 2, p ** k) for x in range(0, p ** k)]
...     for lifted in [set(hensel(r, p, k) for r in roots(s, p ** k))]
...     for actual in [roots(s, p ** (k + 1))]
... )
True