1st version of assignment 2
Correct function