solution for assignment 2 - updated comments