Adding some code for homework 2.2