input: Write a function to check if a number is prime.
  assert is_prime(2) == True
  assert is_prime(3) == True
  assert is_prime(4) == False
  assert is_prime(17) == True
  assert is_prime(20) == False
output: def is_prime(n):
      if n <= 1:
          return False
      if n == 2:
          return True
      if n % 2 == 0:
          return False
      for i in range(3, int(n**0.5) + 1, 2):
          if n % i == 0:
              return False
      return True