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)