2614. Prime In Diagonal

# Time:  precompute: O(MAX_N)
#        runtime: O(n)
# Space: O(MAX_N)

# number theory
def linear_sieve_of_eratosthenes(n):
    primes = []
    spf = [-1]*(n+1)  # the smallest prime factor
    for i in xrange(2, n+1):
        if spf[i] == -1:
            spf[i] = i
            primes.append(i)
        for p in primes:
            if i*p > n or p > spf[i]:
                break
            spf[i*p] = p
    return primes  # len(primes) = O(n/(logn-1)), reference: https://math.stackexchange.com/questions/264544/how-to-find-number-of-prime-numbers-up-to-to-n


MAX_N = 4*10**6
PRIMES = linear_sieve_of_eratosthenes(MAX_N)
PRIMES_SET = set(PRIMES)
class Solution(object):
    def diagonalPrime(self, nums):
        """
        :type nums: List[List[int]]
        :rtype: int
        """
        result = 0
        for i in xrange(len(nums)):
            if nums[i][i] in PRIMES_SET:
                result = max(result, nums[i][i])
            if nums[i][~i] in PRIMES_SET:
                result = max(result, nums[i][~i])
        return result

Solution from kamyu104/LeetCode-Solutions · MIT