Added my code to do the assignment