Integer root algorithm

[BlobCodes]

Feature Request

While building a library, I recently needed a generic iroot method, while the stdlib currently only has an isqrt method.

This is the implementation I created using this article:

  # An algorithm to find the floored `p`th root of `n` (find `x` in term `n=x**p`).
  #
  # Note that while this method uses generics, the given generic `T` must inherit from `Int::Primitive`.
  #
  # ```
  # Math.iroot(64, 3) # => 4 (because 4**3 == 64)
  # ```
  def iroot(n : T, p : T) : T? forall T
    {% T.raise "T must inherit from Int::Primitive" unless T < Int::Primitive %}

    return nil if n < 0 || p <= 0
    return n if n == 0

    n_bits = (sizeof(T) &* 8) &- n.leading_zeros_count
    return T.new(1) if p >= n_bits

    q = (n_bits + p - 1) // p
    x = T.new(1).unsafe_shl(q)
    p_minus_1 = p - 1

    while true
      y = (p_minus_1 * x + n // (x ** p_minus_1)) // p
      return x if x <= y
      x = y
    end
  end

I'd like to ask whether an algorithm like this is considered important enough to be included in the stdlib.

[jzakiya]

See also #8920

[jzakiya]

I played around with this version for iroot and compared it to older version I have.
To get it to work I had to modify it as shown.
I created a template to compare them against (strictly for accuracy and not speed).

def iroot(n, p )
  return nil if n < 0 || p <= 0
  return n if n == 0

  n_bits = n.bit_length
  return typeof(n).new(1) if p >= n_bits

  q = (n_bits + p - 1) // p
  x = typeof(n).new(1) << q
  p_minus_1 = p - 1

  while true
    y = (p_minus_1 &* x &+ n // (x &** p_minus_1)) // p
    return x if x <= y
    x = y
  end
end

module IntRoots

  def inewton(n)        # Newton's method for integer nth root
    return nil if self < 0 && n.even?
    raise "root n not an Integer >= 2" unless n.is_a?(Int) && n > 1
    return self if self == 0 || (self == -1 && n.odd?)
    num = self.abs
    one = typeof(self).new(1)  # value 1 of type self
    e, x = n-1, one << (num.bit_length-1) // n + 1
    while (t = (e &* x &+ num // x &** e) // n) < x
      x = t
    end
    x *= self < 0 ? -1 : 1
  end

  def irootn1(n) # : Int32)
    raise ArgumentError.new "Can't take even root of negative input" if self < 0 && n.even?
    raise ArgumentError.new "Root must be an Integer >= 2" unless n.is_a?(Int) && n > 1
    num = self.abs
    one = typeof(self).new(1)   # value 1 of type self
    root = bitn_mask = one << (num.bit_length - 1) // n
    until (bitn_mask >>= 1) == 0
      root |= bitn_mask
      root ^= bitn_mask if root &** n > num
    end
    root *= (self < 0 ? -1 : 1) # ouput same Integer type as input
  end

  def irootn2(n) #: Int32)
    raise ArgumentError.new "Can't take even root of negative input" if self < 0 && n.even?
    raise ArgumentError.new "Root must be an Integer >= 2" unless n.is_a?(Int) && n > 1
    num = self.abs
    one = typeof(self).new(1)   # value 1 of type self
    root = bitn_mask = one << (b = (num.bit_length - 1) // n)
    numb = 1 << b*n             # root**n
    until (bitn_mask >>= 1) == 0 || numb == num
      root |= bitn_mask
      root ^= bitn_mask if (numb = root &** n) > num
    end
    root *= (self < 0 ? -1 : 1) # ouput same Integer type as input
  end
end

struct Int; include IntRoots end

def display(n, e)
  puts "value for n is: \n#{n}"
  puts "root = #{e}"
  puts "iroot(#{e})";   puts (root = iroot(n, e));   puts "root^#{e}"; puts root.not_nil! &** e; puts "-" * 100
  puts "inewton(#{e})"; puts (root = n.inewton(e));  puts "root^#{e}"; puts root.not_nil! &** e; puts "-" * 100
  puts "irootn1(#{e})"; puts (root = n.irootn1(e));  puts "root^#{e}"; puts root.not_nil! &** e; puts "-" * 100
  puts "irootn2(#{e})"; puts (root = n.irootn2(e));  puts "root^#{e}"; puts root.not_nil! &** e; puts "-" * 100
  puts
end

require "big"

n = 10.to_big_i ** 577
puts
puts "n = #{n}"
e = 2; puts "root = #{e}"
puts "Math.isqrt"; puts (root = Math.isqrt(n));  puts "root^#{e}"; puts root ** e; puts "-" * 100
display(n, 2)
display(n, 3)

n = 123456789012345678
puts
puts "n = #{n}"
e = 2; puts "root = #{e}"
puts "Math.isqrt"; puts (root = Math.isqrt(n));  puts "root^#{e}"; puts root ** e; puts "-" * 100
display(n ,2)
display(n, 3)
display(n, 4)

They all create correct results when n is a BigInt, so I won't show those here.
But inewton and iroot (which is a different implementation of it) give incorrect results for some none BigInt values.
Here's the output to show that for one n.

n = 123456789012345678
root = 2
Math.isqrt
351364182
root^2
123456788392529124
----------------------------------------------------------------------------------------------------
value for n is: 
123456789012345678
root = 2
iroot(2)
351364182
root^2
-2126557980
----------------------------------------------------------------------------------------------------
inewton(2)
351364182
root^2
-2126557980
----------------------------------------------------------------------------------------------------
irootn1(2)
351364182
root^2
123456788392529124
----------------------------------------------------------------------------------------------------
irootn2(2)
351364182
root^2
123456788392529124
----------------------------------------------------------------------------------------------------

value for n is: 
123456789012345678
root = 3
iroot(3)
499236
root^3
1063278144
----------------------------------------------------------------------------------------------------
inewton(3)
499236
root^3
1063278144
----------------------------------------------------------------------------------------------------
irootn1(3)
497933
root^3
123456149902265237
----------------------------------------------------------------------------------------------------
irootn2(3)
497933
root^3
123456149902265237
----------------------------------------------------------------------------------------------------

value for n is: 
123456789012345678
root = 4
iroot(4)
25453
root^4
20964113
----------------------------------------------------------------------------------------------------
inewton(4)
25453
root^4
20964113
----------------------------------------------------------------------------------------------------
irootn1(4)
18744
root^4
123438064202551296
----------------------------------------------------------------------------------------------------
irootn2(4)
18744
root^4
123438064202551296
----------------------------------------------------------------------------------------------------

As you can see irootn1 and irootn2 are more accurate because they don't these type of x &** e operations.
They instead do right shifts >> so the internal values always get smaller.
This is the same type of implementation done by Math.isqrt, which should eliminate arithmetic overflow possibilities.
As a hack for others, do a conditional n.to_big_i when it becomes too big, or switch to gmp version (whichever faster).
When I get time I'll see if I can find|implement an optimum bit-shifting version, similar to Math.iqsrt

Finally, I had to do all these puts root.not_nil! &** e except for output from Math.isqrt, otherwise it gave this error:

 88 | puts "iroot(#{e})";   puts (root = iroot(n, e));   puts "root^#{e}"; puts root &** e; puts "-" * 100
                                                                                     ^--
Error: undefined method '&**' for Nil (compile-time type is (BigInt | Nil))

[HertzDevil]

BigInt is supported via:

lib LibGMP
  fun root = __gmpz_root(rop : MPZ*, op : MPZ*, n : ULong) : Int
end

def iroot(n : BigInt, p : Int::Primitive)
  # error handling / range check not shown
  BigInt.new { |mpz| LibGMP.root(mpz, n, p) }
end

iroot(15625.to_big_i, 6) # => 5.to_big_i
iroot(2187.to_big_i, 7)  # => 3.to_big_i

Invalid arguments should raise instead of returning nil.

[jzakiya]

Invalid arguments should raise instead of returning nil.
Yes, when I got home last night I removed those conditional nil places in those methods to Argument Errors.

I like this version better.

lib LibGMP
  fun root = __gmpz_root(rop : MPZ*, op : MPZ*, n : ULong) : Int
end

def iroot(n, p : Int::Primitive)
  # error handling / range check not shown
  root = BigInt.new { |mpz| LibGMP.root(mpz, n.to_big_i, p) }
  typeof(n).new(root)
end

This works for all Int types except (U||)nt128 because currently no to_(u|i)128 method for BigInts.
Because if doing math with, say, UInt64s all its roots fit in same type, which is same for all types.
If you remove the typeof(n).new(root) then all the root values will be BigInts.

In doing some cursory timings (not benchmarks), this seems just as fast for Math.isqrt too.
So with the proviso to make the output nth roots be same types as the inputs, this could work for all types.

[jzakiya]

This hack (which is not too ugly) will get around the lack of to_(u|i)128 methods for BigInts, and works for all Ints.
Now root will be same type as n.

lib LibGMP
  fun root = __gmpz_root(rop : MPZ*, op : MPZ*, n : ULong) : Int
end

def iroot(n, p : Int::Primitive)
  # error handling / range check not shown
  root = BigInt.new { |mpz| LibGMP.root(mpz, n.to_big_i, p) }
  typeof(n).new(root.to_s)
end