Completed assignment

Made necessary changes to make function work. I also added documenting
comments.