Updated with code for Assignment 2