素数判断 Python

代码

import math
def judge(num):
    if num<2:
        return 0
    for i in range(2,int(math.sqrt(num)+1)):
        if num%i==0:
            return 0
    return 1
num=int(input())
if judge(num)==1:
    print('YES')
else:
    print('NO')