Updated code per assignment