Update code for assignment #2