Code for second assignment