How to check if the number can be represented prime power (nth root is prime or not)

Question

I am trying this problem for a while but getting wrong answer again and again. number can be very large <=2^2014. 22086. Prime Power Test

Explanation about my algorithm:

1. For a Given number I am checking if the number can be represented as form of prime power or not.
2. So the the maximum limit to check for prime power is log n base 2.
3. Finally problem reduced to finding nth root of a number and if it is prime we have our answer else check for all i till log (n base 2) and exit.
4. I have used all sort of optimizations and have tested enormous test-cases and for all my algorithm gives correct answer
5. but Judge says wrong answer.
6. Spoj have another similar problem with small constraints n<=10^18 for which I already got accepted with Python and C++(Best solver in c++)

Here is My python code Please suggest me if I am doing something wrong I am not very proficient in python so my algorithm is a bit lengthy. Thanks in advance.

My Algorithm:

import math
import sys
import fractions
import random
import decimal
write = sys.stdout.write
def sieve(n):
    sqrtn = int(n**0.5)
    sieve = [True] * (n+1)
    sieve[0] = False
    sieve[1] = False
    for i in range(2, sqrtn+1):
        if sieve[i]:
            m = n//i - i
            sieve[i*i:n+1:i] = [False] * (m+1)
    return sieve
def gcd(a, b):
    while b:
        a, b = b, a%b
    return a
def mr_pass(a, s, d, n):
    a_to_power = pow(a, d, n)
    if a_to_power == 1:
        return True
    for i in range(s-1):
        if a_to_power == n - 1:
            return True
        a_to_power = (a_to_power * a_to_power) % n
    return a_to_power == n - 1
isprime=sieve(1000000)
sprime= [2,3,5,7,11,13,17,19,23,29,31,37,41,43,47,53,59,61,67,71,73,79,83,89,97,101,103,107,109,113,127,131,137,139,149,151,157,163,167,173,179,181,191,193,197,199,211,223,227,229,233,239,241,251,257,263,269,271,277,281,283,293,307,311,313,317,331,337,347,349,353,359,367,373,379,383,389,397,401,409,419,421,431,433,439,443,449,457,461,463,467,479,487,491,499,503,509,521,523,541,547,557,563,569,571,577,587,593,599,601,607,613,617,619,631,641,643,647,653,659,661,673,677,683,691,701,709,719,727,733,739,743,751,757,761,769,773,787,797,809,811,821,823,827,829,839,853,857,859,863,877,881,883,887,907,911,919,929,937,941,947,953,967,971,977,983,991,997]
def smooth_num(n):
    c=0
    for a in sprime:
        if(n%a==0):
            c+=1
        if(c>=2):
            return True;
    return False
def is_prime(n):
    if(n<1000000):
        return isprime[n]
    if any((n % p) == 0 for p in sprime):
        return False
    if n==2:
        return True
    d = n - 1
    s = 0
    while d % 2 == 0:
        d >>= 1
        s += 1
    for repeat in range(10):
        a=random.randint(1,n-1)
        if not mr_pass(a, s, d, n):
            return False
    return True
def iroot(n,k):
    hi = 1
    while pow(hi, k) < n:
        hi *= 2
    lo = hi // 2
    while hi - lo > 1:
        mid = (lo + hi) // 2
        midToK = (mid**k)
        if midToK < n:
            lo = mid
        elif n < midToK:
            hi = mid
        else:
            return mid
    if (hi**k) == n:
        return hi
    else:
        return lo
def isqrt(x):
    n = int(x)
    if n == 0:
        return 0
    a, b = divmod(n.bit_length(), 2)
    x = pow(2,(a+b))
    while True:
        y = (x + n//x)>>1
        if y >= x:
            return x
        x = y
maxx=2**1024;minn=2**64
def nth_rootp(n,k):
    return int(round(math.exp(math.log(n)/k),0))
def main():
    for cs in range(int(input())):
        n=int(sys.stdin.readline().strip())
        if(smooth_num(n)):
            write("Invalid order\n")
            continue;
        order = 0;m=0
        power =int(math.log(n,2))
        for i in range(1,power+1):
            if(n<=maxx):
                if i==1:m=n
                elif(i==2):m=isqrt(n)
                elif(i==4):m=isqrt(isqrt(n))
                elif(i==8):m=isqrt(isqrt(isqrt(n)))
                elif(i==16):m=isqrt(isqrt(isqrt(isqrt(n))))
                elif(i==32):m=isqrt(isqrt(isqrt(isqrt(isqrt(n)))))
                elif(i==64):m=isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(n))))))
                elif(i==128):m=isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(n)))))))
                elif(i==256):m=isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(n))))))))
                else:m=int(nth_rootp(n,i))
            else:
                if i==1:m=n
                elif i==2:m=isqrt(n)
                elif(i==4):m=isqrt(isqrt(n))
                elif(i==8):m=isqrt(isqrt(isqrt(n)))
                elif(i==16):m=isqrt(isqrt(isqrt(isqrt(n))))
                elif(i==32):m=isqrt(isqrt(isqrt(isqrt(isqrt(n)))))
                elif(i==64):m=isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(n))))))
                elif(i==128):m=isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(n)))))))
                elif(i==256):m=isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(isqrt(n))))))))
                else:m=iroot(n,i)
            if m<2:
                order=0
                break
            if(is_prime(m) and n==(m**i)):
                write("%d %d\n"%(m,i))
                order = 1
                break
        if(order==0):
            write("Invalid order\n")
main()

Answer 1

from sympy.ntheory import factorint

q=int(input("Give me the number q=")) 

fact=factorint(q)        #We factor the number q=p_1^{n_1}*p_2^{n_2}*...

p_1=list(fact.keys())    #We create a list from keys to be the the numbers p_1,p_2,...

n_1=list(fact.values())  #We create a list from values to be the the numbers n_1,n_2,...

p=int(p_1[0])            

n=int(n_1[0])            

if q!=p**n:              #Check if the number q=p_{1}[0]**n_{1}[0]=p**n.

    print("The number "+str(q)+" is not a prime power")

else:

    print("The number "+str(q)+" is a prime power")

    print("The prime number p="+str(p))

    print("The natural number n="+str(n))

Answer 2

your code look like a little overcomplicated for this task, I will not bother to check it, but the thing you need are the following

1. is_prime, naturally
2. a prime generator, optional
3. calculate the nth root of a number in a precise way

for the first one I recommend the deterministic form of the Miller-Rabin test with a appropriate set of witness to guaranty a exact result until 1543267864443420616877677640751301 (1.543 x 10³³) for even bigger numbers you can use the probabilistic one or use a bigger list of witness chosen at your criteria

with all that a template for the solution is as follow

import math

def is_prime(n):
    ...

def sieve(n):
    "list of all primes p such that p<n"
    ...

def inthroot(x,n):
    "calculate floor(x**(1/n))"
    ...

def is_a_power(n):
    "return (a,b) if n=a**b otherwise throw ValueError"
    for b in sieve( math.log2(n) +1 ):
        a = inthroot(n,b)
        if a**b == n:
            return a,b
    raise ValueError("is not a power")

def smooth_factorization(n):
    "return (p,e) where p is prime and n = p**e if such value exists, otherwise throw ValueError"
    e=1
    p=n
    while True:
        try:
            p,n = is_a_power(p)
            e = e*n
        except ValueError:
            break
    if is_prime(p): 
        return p,e
    raise ValueError

def main():
    for test in range( int(input()) ):
        try:
            p,e = smooth_factorization( int(input()) )
            print(p,e)
        except ValueError:
            print("Invalid order")

main()

And the code above should be self explanatory

Filling the blacks

As you are familiar with Miller-Rabin test, I will only mention that if you are interested you can find a implementation of the determinist version here just update the list of witness and you are ready to go.

For the sieve, just change the one you are using to return a list with primes number like this for instance [ p for p,is_p in enumerate(sieve) if is_p ]

With those out of the way, the only thing left is calculate the nth root of the number and to do that in a precise way we need to get rip of that pesky floating point arithmetic that only produce headaches, and the answer is implement the Nth root algorithm using only integer arithmetic, which is pretty similar to the one of isqrt that you already use, I guide myself with the one made by Mark Dickinson for cube root and generalize it and I get this

def inthroot(A, n) :
    "calculate floor( A**(1/n) )"
    #https://en.wikipedia.org/wiki/Nth_root_algorithm
    #https://en.wikipedia.org/wiki/Nth_root#nth_root_algorithm
    #https://stackoverflow.com/questions/35254566/wrong-answer-in-spoj-cubert/35276426#35276426
    #https://stackoverflow.com/questions/39560902/imprecise-results-of-logarithm-and-power-functions-in-python/39561633#39561633
    if A<0:
        if n%2 == 0:
            raise ValueError
        return - inthroot(-A,n)
    if A==0:
        return 0
    n1 = n-1
    if A.bit_length() < 1024: # float(n) safe from overflow
        xk = int( round( pow(A,1.0/n) ) )
        xk = ( n1*xk + A//pow(xk,n1) )//n # Ensure xk >= floor(nthroot(A)).
    else:
        xk = 1 << -(-A.bit_length()//n) # 1 << sum(divmod(A.bit_length(),n))
                                        # power of 2 closer but greater than the nth root of A
    while True:
        sig = A // pow(xk,n1)
        if xk <= sig:
            return xk
        xk = ( n1*xk + sig )//n

and with all the above you can solve the problem without inconvenient

Answer 3

I'm not going to read all that code, though I suspect the problem is floating-point inaccuracy. Here is my program to determine if a number n is a prime power; it returns the prime p and the power k:

# prime power predicate

from random import randint
from fractions import gcd

def findWitness(n, k=5): # miller-rabin
    s, d = 0, n-1
    while d % 2 == 0:
        s, d = s+1, d/2
    for i in range(k):
        a = randint(2, n-1)
        x = pow(a, d, n)
        if x == 1 or x == n-1: continue
        for r in range(1, s):
            x = (x * x) % n
            if x == 1: return a
            if x == n-1: break
        else: return a
    return 0

# returns p,k such that n=p**k, or 0,0
# assumes n is an integer greater than 1
def primePower(n):
    def checkP(n, p):
        k = 0
        while n > 1 and n % p == 0:
            n, k = n / p, k + 1
        if n == 1: return p, k
        else: return 0, 0
    if n % 2 == 0: return checkP(n, 2)
    q = n
    while True:
        a = findWitness(q)
        if a == 0: return checkP(n, q)
        d = gcd(pow(a,q,n)-a, q)
        if d == 1 or d == q: return 0, 0
        q = d

The program uses Fermat's Little Theorem and exploits the witness a to the compositeness of n that is found by the Miller-Rabin algorithm. It is given as Algorithm 1.7.5 in Henri Cohen's book A Course in Computational Algebraic Number Theory. You can see the program in action at http://ideone.com/cNzQYr.

Answer 4

this is not really an answer, but I don't have enough space to write it as a comment.

So, if the problem still not solved, you may try the following function for nth_rootp, though it is a bit ugly (it is just a binary search to find the precise value of the function):

def nth_rootp(n,k):
    r = int(round(math.log(n,2)/k))
    left = 2**(r-1)
    right = 2**(r+1)
    if left**k == n:
        return left
    if right**k == n:
        return right
    while left**k < n and right**k > n:
        tmp = (left + right)/2
        if tmp**k == n:
            return tmp
        if tmp == left or tmp == right:
            return tmp
        if tmp**k < n:
            left = tmp
        else:
            if tmp**k > n:
                right = tmp