python/Code_Challenge_medium_1.py

53 lines
1.6 KiB
Python

'''
We know that prime numbers are positive integers that have exactly two distinct positive divisors. Similarly, we define a positive integer t as a T-prime
if it has exactly three distinct positive divisors.
Given an array of positive integers, return an array of True or False for each integer if it T-prime
Example 1
Input:
[4, 5, 6]
Output:
[True, False, False]
'''
from typing import List
class Solution(object):
def solve(self, nums:List[int]) -> List[bool]:
"""
:type nums: List[int]
:rtype: List[bool]
"""
def is_t_prime(n):
if n < 4:
return False
root = int(n**0.5)
return root * root == n and self.is_prime(root)
self.is_prime = self.sieve_of_eratosthenes(max(nums))
return [is_t_prime(num) for num in nums]
def sieve_of_eratosthenes(self, n):
"""
Returns a function that checks if a number is prime using the Sieve of Eratosthenes.
"""
is_prime = [True] * (n + 1)
is_prime[0] = is_prime[1] = False
for i in range(2, int(n**0.5) + 1):
if is_prime[i]:
for j in range(i * i, n + 1, i):
is_prime[j] = False
return lambda x: is_prime[x] if x <= n else False
# Example usage:
if __name__ == "__main__":
solution = Solution()
print(solution.solve([4, 5, 6])) # Output: [True, False, False]
print(solution.solve([9, 25, 49])) # Output: [True, True, True]
print(solution.solve([1, 2, 3, 10])) # Output: [False, False, False, False]