add updated code for assignment