Prime Number Generating Puzzle (Edge Cases)

Question

I'm solving the following problem on a coding site. It's failing for some edge cases on the tests (hidden tests), but I'm not sure what they are. Anyone see any issues with this?

Problem: Let A be a string of all the prime numbers sequentially squashed together (i.e. 235711131719...). Given an index n, return a string of 5 digits where the first digit is at index n in A.

e.g. foo(0) => 23571 and foo(10) => 19232

Here's my code:

def gen_primes():                                                                                                                                                                                                    
    A = {}                                                                                                                                                                                                           
    i = 2                                                                                                                                                                                                            
    while True:                                                                                                                                                                                                      
        if i not in A:                                                                                                                                                                                               
            yield i                                                                                                                                                                                                  
            A[i * i] = [i]                                                                                                                                                                                           
        else:                                                                                                                                                                                                        
            for p in A[i]:                                                                                                                                                                                           
                A.setdefault(p + i, []).append(p)                                                                                                                                                                    
            del A[i]                                                                                                                                                                                                 
        i += 1                                                                                                                                                                                                       

def answer(n):                                                                                                                                                                                                       

    counter = 0                                                                                                                                                                                                      
    prime_string = ""                                                                                                                                                                                                

    for p in gen_primes():                                                                                                                                                                                           
        if (counter >= n):
            prime_string += str(p)                                                                                                                                                                                   
        counter += len(str(p))                                                                                                                                                                                       
        if len(prime_string) >= 5:                                                                                                                                                                                   
            break                                                                                                                                                                                                    

    return prime_string[:5]    

Answer 1

I think this could break for primes with more than one digit:

Let's assume that we have arrived at primes with three digits, like 103. Counter is 10 and n is 11 (this is just an example, I don't know if this exact constellation would turn up)

Then we would need to use the digits "03" out of "103". But since counter is smaller than n, the entire prime is skipped. The program would continue with 107.

You could fix this by removing counter entirely: Always add primes to the string, break out of the loop if the length of the string is n+5 or more.

EDIT:

I've checked your code: An example is answer(5) and answer(6). With your code, both calls result in "13171". The second digit of "11" is skipped.

With this code:

def answer(n):                                                                                                                                                                                                       

    counter = 0                                                                                                                                                                                                      
    prime_string = ""                                                                                                                                                                                                

    for p in gen_primes():
        prime_string += str(p)                                                                                                                                                                                     
        if len(prime_string) >= n+5:                                                                                                                                                                                   
            break                                                                                                                                                                                                    

    return prime_string[n:n+5] 

They result in

11317  # answer(5)
13171  # answer(6)

Answer 2

Here's my answer; it is similar to that of @lhk, but uses a sliding window instead of storing the entire string:

phil@haydn ~ $ python
Python 2.7.5+ (default, Sep 17 2013, 15:31:50) 
[GCC 4.8.1] on linux2
Type "help", "copyright", "credits" or "license" for more information.
>>> def primegen(): # http://stackoverflow.com/a/20660551
...     yield 2; yield 3          # prime (!) the pump
...     ps = primegen()           # sieving primes
...     p = next(ps) and next(ps) # first sieving prime
...     D, q, c = {}, p*p, p      # initialize
...     def add(x, s):            # insert multiple/stride
...         while x in D: x += s  #   find unused multiple
...         D[x] = s              #   save multiple/stride
...     while True:               # infinite list
...         c += 2                # next odd candidate
...         if c in D:            # c is composite
...             s = D.pop(c)      #   fetch stride
...             add(c+s, s)       #   add next multiple
...         elif c < q: yield c   # c is prime; yield it
...         else: # (c == q)      # add sqrt(c) to sieve
...             add(c+p+p, p+p)   #   insert in sieve
...             p = next(ps)      #   next sieving prime
...             q = p * p         #   ... and its square
... 
>>> def primeSubstr(n): # nth character in concatenated primes
...     ps = primegen()           # sequence of prime numbers
...     i, s = 0, ''              # current index, sliding window
...     while True:               # do ... until i == n
...         if len(s) < 5:        # sliding window too small
...             s+=str(next(ps))  #   add next prime to window
...         elif i < n:           # not yet to target index
...             i, s = i+1, s[1:] #   slide window to right
...         else: return s[:5]    # return desired substring
... 
>>> primeSubstr(0)
'23571'
>>> primeSubstr(5)
'11317'
>>> primeSubstr(10)
'19232'
>>> primeSubstr(15)
'93137'
>>> primeSubstr(20)
'41434'
>>> primeSubstr(1000)
'98719'
>>> CTRL-D

I will discuss this problem tomorrow at my blog.

Answer 3

this looks like a job for itertools, how about this

>>> from sympy import sieve
>>> from itertools import islice, chain
>>> def primeSubStr(n):
        primes=iter(sieve)
        next(primes)
        return "".join( islice( chain.from_iterable(map(str,primes)), n, n+5))

>>> primeSubStr(0)
'23571'
>>> primeSubStr(10)
'19232'
>>> primeSubStr(5)
'11317'
>>> primeSubStr(15)
'93137'
>>> primeSubStr(20)
'41434'
>>> primeSubStr(1000)
'98719'
>>> primeSubStr(2000)
'98940'
>>> 

For simplicity I use sympy.sieve to get the prime numbers, but any method of getting the primes would work just fine, like the ones here for instance.

Now the interesting part, the map is obvious enough, with that we get a stream of "2", "3", "5", "7", "11", "13", ... "101", "103", ... then with chain.from_iterable that stream is discomposed in a stream of singles digit "2", "3", "5", "7", "1","1", "1","3", ... "1","0","1", "1","0","3", ... and finally with islice we take the part we want and join it for the final result