Updated Code for Assignment