Updated code for assignment 2