Updated code for Assignment 2