Added some code for lecture 4