@ -1476,6 +1476,266 @@ count_set_bits(unsigned long n)
return count ;
}
/* Integer square root
Given a nonnegative integer ` n ` , we want to compute the largest integer
` a ` for which ` a * a < = n ` , or equivalently the integer part of the exact
square root of ` n ` .
We use an adaptive - precision pure - integer version of Newton ' s iteration . Given
a positive integer ` n ` , the algorithm produces at each iteration an integer
approximation ` a ` to the square root of ` n > > s ` for some even integer ` s ` ,
with ` s ` decreasing as the iterations progress . On the final iteration , ` s ` is
zero and we have an approximation to the square root of ` n ` itself .
At every step , the approximation ` a ` is strictly within 1.0 of the true square
root , so we have
( a - 1 ) * * 2 < ( n > > s ) < ( a + 1 ) * * 2
After the final iteration , a check - and - correct step is needed to determine
whether ` a ` or ` a - 1 ` gives the desired integer square root of ` n ` .
The algorithm is remarkable in its simplicity . There ' s no need for a
per - iteration check - and - correct step , and termination is straightforward : the
number of iterations is known in advance ( it ' s exactly ` floor ( log2 ( log2 ( n ) ) ) `
for ` n > 1 ` ) . The only tricky part of the correctness proof is in establishing
that the bound ` ( a - 1 ) * * 2 < ( n > > s ) < ( a + 1 ) * * 2 ` is maintained from one
iteration to the next . A sketch of the proof of this is given below .
In addition to the proof sketch , a formal , computer - verified proof
of correctness ( using Lean ) of an equivalent recursive algorithm can be found
here :
https : / / github . com / mdickinson / snippets / blob / master / proofs / isqrt / src / isqrt . lean
Here ' s Python code equivalent to the C implementation below :
def isqrt ( n ) :
" " "
Return the integer part of the square root of the input .
" " "
n = operator . index ( n )
if n < 0 :
raise ValueError ( " isqrt() argument must be nonnegative " )
if n = = 0 :
return 0
c = ( n . bit_length ( ) - 1 ) / / 2
a = 1
d = 0
for s in reversed ( range ( c . bit_length ( ) ) ) :
e = d
d = c > > s
a = ( a < < d - e - 1 ) + ( n > > 2 * c - e - d + 1 ) / / a
assert ( a - 1 ) * * 2 < n > > 2 * ( c - d ) < ( a + 1 ) * * 2
return a - ( a * a > n )
Sketch of proof of correctness
- - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
The delicate part of the correctness proof is showing that the loop invariant
is preserved from one iteration to the next . That is , just before the line
a = ( a < < d - e - 1 ) + ( n > > 2 * c - e - d + 1 ) / / a
is executed in the above code , we know that
( 1 ) ( a - 1 ) * * 2 < ( n > > 2 * ( c - e ) ) < ( a + 1 ) * * 2.
( since ` e ` is always the value of ` d ` from the previous iteration ) . We must
prove that after that line is executed , we have
( a - 1 ) * * 2 < ( n > > 2 * ( c - d ) ) < ( a + 1 ) * * 2
To faciliate the proof , we make some changes of notation . Write ` m ` for
` n > > 2 * ( c - d ) ` , and write ` b ` for the new value of ` a ` , so
b = ( a < < d - e - 1 ) + ( n > > 2 * c - e - d + 1 ) / / a
or equivalently :
( 2 ) b = ( a < < d - e - 1 ) + ( m > > d - e + 1 ) / / a
Then we can rewrite ( 1 ) as :
( 3 ) ( a - 1 ) * * 2 < ( m > > 2 * ( d - e ) ) < ( a + 1 ) * * 2
and we must show that ( b - 1 ) * * 2 < m < ( b + 1 ) * * 2.
From this point on , we switch to mathematical notation , so ` / ` means exact
division rather than integer division and ` ^ ` is used for exponentiation . We
use the ` √ ` symbol for the exact square root . In ( 3 ) , we can remove the
implicit floor operation to give :
( 4 ) ( a - 1 ) ^ 2 < m / 4 ^ ( d - e ) < ( a + 1 ) ^ 2
Taking square roots throughout ( 4 ) , scaling by ` 2 ^ ( d - e ) ` , and rearranging gives
( 5 ) 0 < = | 2 ^ ( d - e ) a - √ m | < 2 ^ ( d - e )
Squaring and dividing through by ` 2 ^ ( d - e + 1 ) a ` gives
( 6 ) 0 < = 2 ^ ( d - e - 1 ) a + m / ( 2 ^ ( d - e + 1 ) a ) - √ m < 2 ^ ( d - e - 1 ) / a
We ' ll show below that ` 2 ^ ( d - e - 1 ) < = a ` . Given that , we can replace the
right - hand side of ( 6 ) with ` 1 ` , and now replacing the central
term ` m / ( 2 ^ ( d - e + 1 ) a ) ` with its floor in ( 6 ) gives
( 7 ) - 1 < 2 ^ ( d - e - 1 ) a + m / / 2 ^ ( d - e + 1 ) a - √ m < 1
Or equivalently , from ( 2 ) :
( 7 ) - 1 < b - √ m < 1
and rearranging gives that ` ( b - 1 ) ^ 2 < m < ( b + 1 ) ^ 2 ` , which is what we needed
to prove .
We ' re not quite done : we still have to prove the inequality ` 2 ^ ( d - e - 1 ) < =
a ` that was used to get line ( 7 ) above . From the definition of ` c ` , we have
` 4 ^ c < = n ` , which implies
( 8 ) 4 ^ d < = m
also , since ` e = = d > > 1 ` , ` d ` is at most ` 2 e + 1 ` , from which it follows
that ` 2 d - 2 e - 1 < = d ` and hence that
( 9 ) 4 ^ ( 2 d - 2 e - 1 ) < = m
Dividing both sides by ` 4 ^ ( d - e ) ` gives
( 10 ) 4 ^ ( d - e - 1 ) < = m / 4 ^ ( d - e )
But we know from ( 4 ) that ` m / 4 ^ ( d - e ) < ( a + 1 ) ^ 2 ` , hence
( 11 ) 4 ^ ( d - e - 1 ) < ( a + 1 ) ^ 2
Now taking square roots of both sides and observing that both ` 2 ^ ( d - e - 1 ) ` and
` a ` are integers gives ` 2 ^ ( d - e - 1 ) < = a ` , which is what we needed . This
completes the proof sketch .
*/
/*[clinic input]
math . isqrt
n : object
/
Return the integer part of the square root of the input .
[ clinic start generated code ] */
static PyObject *
math_isqrt ( PyObject * module , PyObject * n )
/*[clinic end generated code: output=35a6f7f980beab26 input=5b6e7ae4fa6c43d6]*/
{
int a_too_large , s ;
size_t c , d ;
PyObject * a = NULL , * b ;
n = PyNumber_Index ( n ) ;
if ( n = = NULL ) {
return NULL ;
}
if ( _PyLong_Sign ( n ) < 0 ) {
PyErr_SetString (
PyExc_ValueError ,
" isqrt() argument must be nonnegative " ) ;
goto error ;
}
if ( _PyLong_Sign ( n ) = = 0 ) {
Py_DECREF ( n ) ;
return PyLong_FromLong ( 0 ) ;
}
c = _PyLong_NumBits ( n ) ;
if ( c = = ( size_t ) ( - 1 ) ) {
goto error ;
}
c = ( c - 1U ) / 2U ;
/* s = c.bit_length() */
s = 0 ;
while ( ( c > > s ) > 0 ) {
+ + s ;
}
a = PyLong_FromLong ( 1 ) ;
if ( a = = NULL ) {
goto error ;
}
d = 0 ;
while ( - - s > = 0 ) {
PyObject * q , * shift ;
size_t e = d ;
d = c > > s ;
/* q = (n >> 2*c - e - d + 1) // a */
shift = PyLong_FromSize_t ( 2U * c - d - e + 1U ) ;
if ( shift = = NULL ) {
goto error ;
}
q = PyNumber_Rshift ( n , shift ) ;
Py_DECREF ( shift ) ;
if ( q = = NULL ) {
goto error ;
}
Py_SETREF ( q , PyNumber_FloorDivide ( q , a ) ) ;
if ( q = = NULL ) {
goto error ;
}
/* a = (a << d - 1 - e) + q */
shift = PyLong_FromSize_t ( d - 1U - e ) ;
if ( shift = = NULL ) {
Py_DECREF ( q ) ;
goto error ;
}
Py_SETREF ( a , PyNumber_Lshift ( a , shift ) ) ;
Py_DECREF ( shift ) ;
if ( a = = NULL ) {
Py_DECREF ( q ) ;
goto error ;
}
Py_SETREF ( a , PyNumber_Add ( a , q ) ) ;
Py_DECREF ( q ) ;
if ( a = = NULL ) {
goto error ;
}
}
/* The correct result is either a or a - 1. Figure out which, and
decrement a if necessary . */
/* a_too_large = n < a * a */
b = PyNumber_Multiply ( a , a ) ;
if ( b = = NULL ) {
goto error ;
}
a_too_large = PyObject_RichCompareBool ( n , b , Py_LT ) ;
Py_DECREF ( b ) ;
if ( a_too_large = = - 1 ) {
goto error ;
}
if ( a_too_large ) {
Py_SETREF ( a , PyNumber_Subtract ( a , _PyLong_One ) ) ;
}
Py_DECREF ( n ) ;
return a ;
error :
Py_XDECREF ( a ) ;
Py_DECREF ( n ) ;
return NULL ;
}
/* Divide-and-conquer factorial algorithm
*
* Based on the formula and pseudo - code provided at :
@ -2737,6 +2997,7 @@ static PyMethodDef math_methods[] = {
MATH_ISFINITE_METHODDEF
MATH_ISINF_METHODDEF
MATH_ISNAN_METHODDEF
MATH_ISQRT_METHODDEF
MATH_LDEXP_METHODDEF
{ " lgamma " , math_lgamma , METH_O , math_lgamma_doc } ,
MATH_LOG_METHODDEF
Mark Dickinson
7 years ago
GitHub