HackerRank - Sam's Prime Sum

gistfile1.txt

import sys
def find_index(arr, n):
    maximum_element = max(arr)    #3
    global prime
    #prime = [i for i in range(maximum_element + 1)]
    if len(prime) < maximum_element+1:   #2>4
        difference = (maximum_element +1) - len(prime) #4-2 = 2
        prime += [True] * difference
        p = (maximum_element +1) - len(prime)
        #prime[0] = False
        #prime[1] = False
        while p * p < maximum_element:
            if prime[p]:
                for i in range(p * p, maximum_element + 1, p):
                    prime[i] = False
            p += 1

    prime_sum_from_left = 0
    prime_sum_from_right = 0

    first_array = [0] * n
    second_array = [0] * n
    for i in range(n):

        first_array[i] = prime_sum_from_left
        second_array[n - i -1 ] = prime_sum_from_right

        if prime[arr[i]]:
            prime_sum_from_left += arr[i]
        if prime[arr[n - i - 1]]:
            prime_sum_from_right += arr[n - i - 1]


    for i in range(n):
       if first_array[i] == second_array[i]:
          return i
    # No index is found.
    return -1


cases = int(input())
prime = [False,False]
mainArr = []
for x in range(cases):
    arrLen = int(input())
    tempArr = list(map(int,sys.stdin.readline().split()))
    mainArr.append(tempArr)
for x in mainArr:
    length = len(x)
    answer = str(find_index(x,length))
    sys.stdout.write(answer + '\n')