Last week, we wrote some code to test whether or not a given number was prime.
n = 132
n_is_prime = True
for d in range(2,n):
if n % d == 0:
n_is_prime = False
break
if n_is_prime:
print('{} is a prime number'.format(n))
else:
print('{} is not a prime number'.format(n))
132 is not a prime number
Very often, we want to write code that can be applied to many different inputs. We can accomplish this by writing a Python function.
The syntax for defining functions in Python is:
def <function name>(<some inputs>):
<do something>def this_is_a_function():
print('Hello world!')
We can call a function after defining it:
this_is_a_function()
Hello world!
def function_with_one_input(n):
if n < 100:
print('{} is a small number'.format(n))
else:
print('{} is a large number'.format(n))
function_with_one_input(50)
50 is a small number
function_with_one_input(150)
150 is a large number
Exercise: Write a function is_prime that takes in a variable n and prints whether or not n is prime.
def is_prime(n):
n_is_prime = True
for d in range(2,n):
if n % d == 0:
n_is_prime = False
break
if n_is_prime:
print('{} is a prime number'.format(n))
else:
print('{} is not a prime number'.format(n))
is_prime(133)
133 is not a prime number
It is often more useful to have functions like the above return either a True or a False Boolean depending on whether n is prime or not.
More generally, we often want functions to return some value to the caller. This can be done using a return statement:
my_string = 'Hello World!'
def is_prime(n):
n_is_prime = True
print(my_string)
for d in range(2,n):
if n % d == 0:
n_is_prime = False
break
return n_is_prime
is_prime(127)
Hello World!
True
is_prime(131)
Hello World!
True
if n_is_prime:
print('{} is a prime number'.format(n))
else:
print('{} is not a prime number'.format(n))
132 is not a prime number
Note: Variables that are defined within a function only exist within that function. In the above example, we've defined n_is_prime within the is_prime function. This is a completely different variable to the globally defined variable n_is_prime that was defined at the top of this notebook. If a function tries to use a variable name that it has defined, it will look outside the function for that variable name.
Note: A function will terminate as soon as it hits a return statement.
def this_is_a_function(a,b):
print('Hello world!')
return a * b
print('Goodbye!')
this_is_a_function(5,7)
Hello world!
35
def is_prime(n):
for d in range(2,n):
if n % d == 0:
return False
return True
Exercise: Write a function my_primes that takes in an input n and returns a list of all prime numbers between 2 and n (inclusively).
def my_primes(n):
primes = []
for i in range(2,n+1):
if is_prime(i):
primes.append(i)
return primes
my_primes(29)
[2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
How efficient is our is_prime function? It would helpful if we could time our code to compare efficiency. We can use the time module to do this.
Many useful functions are contained modules but are not immediately accessible in a new notebook. We can use the import function to gain access to a module.
import time
time.time()
1738616063.7513962
Let's use this to time our my_primes function:
t0 = time.time()
primes = my_primes(100000)
t1 = time.time()
print(t1 - t0)
20.80789303779602
There are other ways to get access to objects inside a module. For example, the syntax import <object inside module> from <module> will pull the object outside of the module for normal use.
from time import time_ns
time_ns()
1738616301957215100
time.time_ns()
1738616318741235800
time.asctime()
'Mon Feb 3 15:59:05 2025'
asctime()
--------------------------------------------------------------------------- NameError Traceback (most recent call last) ~\AppData\Local\Temp\ipykernel_4468\3577197693.py in <module> ----> 1 asctime() NameError: name 'asctime' is not defined
from time import asctime
asctime()
'Mon Feb 3 15:59:43 2025'
Let's look again at our is_prime function:
def is_prime(n):
for d in range(2,n):
if n % d == 0:
return False
return True
From our discussion, a number $n$ can only be non-prime if it has divisor $d$ less than or equal to $\sqrt{n}$.
We can pull the sqrt function from the math module to compute square roots:
from math import sqrt
def is_prime(n):
for d in range(2,sqrt(n)):
if n % d == 0:
return False
return True
is_prime(29)
--------------------------------------------------------------------------- TypeError Traceback (most recent call last) ~\AppData\Local\Temp\ipykernel_4468\3856216990.py in <module> ----> 1 is_prime(29) ~\AppData\Local\Temp\ipykernel_4468\4195421119.py in is_prime(n) 1 def is_prime(n): ----> 2 for d in range(2,sqrt(n)): 3 if n % d == 0: 4 return False 5 return True TypeError: 'float' object cannot be interpreted as an integer
The range function only accepts integer inputs, so let's convert sqrt(n) to an integer type:
def is_prime(n):
for d in range(2,int(sqrt(n))+1):
if n % d == 0:
return False
return True
is_prime(5)
True
Let's compare the speed between this new version and our old version:
def is_prime_old(n):
for d in range(2,n):
if n % d == 0:
return False
return True
def is_prime(n):
for d in range(2,int(sqrt(n))+1):
if n % d == 0:
return False
return True
n = 100000007
t0 = time.time()
is_prime_old(n)
t1 = time.time()
print('Unoptimized time:', t1 - t0)
t0 = time.time()
is_prime(n)
t1 = time.time()
print('Optimized time:', t1 - t0)
Unoptimized time: 5.047978639602661 Optimized time: 0.0
Let's compare the use of these two versions in the my_primes function:
def my_primes_old(n):
primes = []
for i in range(2,n+1):
if is_prime_old(i):
primes.append(i)
return primes
def my_primes(n):
primes = []
for i in range(2,n+1):
if is_prime(i):
primes.append(i)
return primes
n = 100000
t0 = time.time()
my_primes_old(n)
t1 = time.time()
print('Unoptimized time:', t1 - t0)
t0 = time.time()
my_primes(n)
t1 = time.time()
print('Optimized time:', t1 - t0)
Unoptimized time: 21.64876437187195 Optimized time: 0.1591646671295166
Exercise: Write a function that takes in an integer n and returns a list of prime numbers that divide n. (We will need to modify this function to account multiplicity of prime factors to create the primary function for the project).
Exercise: Write a function is_prime_like that takes in an integer n and returns True if n is prime-like and False otherwise.
Reminder: A number $n$ is prime-like if:
$$ a^n \equiv a \pmod n $$for all integers $0 \leq a < n$.
def is_prime_like(n):
for a in range(n):
if a**n % n == a % n:
continue
else:
return False
return True
In the above function, we used a if statement where we only wanted to do something if the condition was False. We can use a not operator to invert a Boolean expression:
def is_prime_like(n):
for a in range(n):
if not a**n % n == a % n:
return False
return True
is_prime_like(5)
True
is_prime_like(4)
False
is_prime_like(100000007)
n = 100000007
a = n - 1
pow(a,n,n)
100000006
not, and, or¶We can invert and/or combine Boolean expressions using the operators not, and, and or:
from math import sqrt
def is_prime(n):
for d in range(2,int(sqrt(n))+1):
if n % d == 0:
return False
return True
Suppose we want to find twin primes, that is, pairs of primes that differ by 2.
twin_primes = []
for n in range(2,10000):
if is_prime(n) and is_prime(n+2):
twin_primes.append([n,n+2])
#twin_primes
Find all numbers that are either a perfect square or a perfect cube up to 10,000:
squares_or_cubes = []
for n in range(10000):
sqrtn = sqrt(n)
cubertn = n**(1/3)
if sqrtn == int(sqrtn) or cubertn == int(cubertn):
squares_or_cubes.append(n)
#squares_or_cubes
See Project report guide on class webpage.
As an example, suppose we want to explore the topic of twin primes.
from math import pi
def cone_volume(r,h):
return pi * r**2 * h / 3
In general, the order in which we enter arguments to a function matters. For example, cone_volume(2,5) finds the volume of a cone with r=2 and h=5, while cone_volume(5,2) finds the volume with r=5 and h=2.
cone_volume(2,5)
20.943951023931955
cone_volume(5,2)
52.35987755982989
In Python, we can alternatively enter arguments in any order, as long as we label them:
cone_volume(r=2, h=5)
20.943951023931955
cone_volume(h=5, r=2)
20.943951023931955
cone_volume(2, h=5)
20.943951023931955
These are called keyword arguments, that is, when we enter <variable name>=<value> as an argument.
There are also different ways of defining input variables for a function. In particular, we can assign default values to inputs that we want to be optional.
For example, suppose we want the height h to be optional, with a default value of h=5.
from math import pi
def cone_volume(r,h=5):
return pi * r**2 * h / 3
cone_volume(5,2)
52.35987755982989
cone_volume(2)
20.943951023931955
Arguments that are not assigned default values are called positional arguments. Positional arguments must always come before keyword arguments (when calling a function) and arguments with default values (when defining a function).
import math
def cone_volume(r,h=5, pi=math.pi):
return pi * r**2 * h / 3
cone_volume(2, pi=3)
20.0