Digit-by-digit square root in Python
Computes square root one digit at a time in the manner of long division.
- File
digit_by_digit_square_root.py, 56 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/digit-by-digit-square-root/. The authors are listed in prov/SH6MP2.jsonl.
def digit_sqrt(a: list[int], base: int) -> tuple[list[int], int]:
"""Integer square root of a number given as base-`base` digits.
Precondition: base >= 2 and 0 <= a[i] < base, digits least significant
first (an empty list is 0). Returns (root digits, remainder) with
root = isqrt(A), remainder = A - root^2 and 0 <= remainder <= 2 * root;
the root has ceil(len(a) / 2) digits, least significant first. The
remainder r and the partial root q are Python ints of O(n) digits. Cost:
n steps for a root of n digits, each with O(log base) trial products of
an O(n)-digit number by one digit, so Theta(n^2) digit operations for a
fixed base.
"""
if base < 2 or any(not 0 <= v < base for v in a):
raise ValueError("digits must lie in 0 .. base - 1")
digits = list(a)
if len(digits) % 2:
digits.append(0)
n = len(digits) // 2
q, r = 0, 0 # @line L1
root: list[int] = []
for k in range(n - 1, -1, -1): # @line L2
r = r * base * base + digits[2 * k + 1] * base + digits[2 * k] # @line L3
lo, hi = 0, base - 1 # @line L4: largest x with (2 q base + x) x <= r
while lo < hi:
mid = (lo + hi + 1) // 2
if (2 * q * base + mid) * mid <= r:
lo = mid
else:
hi = mid - 1
x = lo
r -= (2 * q * base + x) * x # @line L5
q = q * base + x # @line L6
root.append(x)
root.reverse()
return root, r # @line L7
def to_digits(x: int, base: int) -> list[int]:
"""The digits of x >= 0 in the given base, least significant first."""
out: list[int] = []
while x:
x, d = divmod(x, base)
out.append(d)
return out
def from_digits(digits: list[int], base: int) -> int:
"""The value of a digit list, least significant first."""
value = 0
for d in reversed(digits):
value = value * base + d
return value
test_digit_by_digit_square_root.py 72 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/digit-by-digit-square-root/. The authors are listed in prov/SH6MP2.jsonl.
import math
import random
from digit_by_digit_square_root import digit_sqrt, from_digits, to_digits
def check(x: int, base: int) -> None:
root_digits, r = digit_sqrt(to_digits(x, base), base)
q = from_digits(root_digits, base)
assert q == math.isqrt(x)
assert r == x - q * q
assert 0 <= r <= 2 * q
assert all(0 <= d < base for d in root_digits)
def test_example() -> None:
root, r = digit_sqrt(to_digits(152399025, 10), 10)
assert root == [5, 4, 3, 2, 1] and r == 0
root, r = digit_sqrt(to_digits(200, 2), 2)
assert from_digits(root, 2) == 14 and r == 4
def test_edges() -> None:
for base in (2, 3, 10, 16, 256):
assert digit_sqrt([], base) == ([], 0)
for x in (0, 1, 2, 3, 4, base - 1, base, base * base - 1, base * base):
check(x, base)
for length in (1, 2, 3, 7, 8):
check(base**length - 1, base) # all digits base - 1
check(base**length, base)
def test_squares_and_neighbours() -> None:
rng = random.Random(77)
for base in (2, 10, 16, 256):
for _ in range(60):
q = rng.randrange(1, 1 << rng.randint(1, 120))
for x in (q * q - 1, q * q, q * q + 2 * q, (q + 1) ** 2):
check(x, base)
def test_float_trap() -> None:
a = (2**27 + 1) ** 2 - 1
root, r = digit_sqrt(to_digits(a, 16), 16)
assert from_digits(root, 16) == 2**27 and r == 2**28
def test_exhaustive_small() -> None:
for base in (2, 3, 4, 10):
for x in range(2000):
check(x, base)
def test_random_against_isqrt() -> None:
rng = random.Random(9001)
for _ in range(600):
base = rng.choice((2, 3, 7, 10, 16, 256))
x = rng.randrange(1 << rng.randint(1, 400))
check(x, base)
def test_bad_digits() -> None:
for a, base in (([2], 2), ([0, -1], 10), ([1], 1)):
try:
digit_sqrt(a, base)
except ValueError:
pass
else:
raise AssertionError("expected ValueError")
What the next level needs
- the tests did not run in this build