Adding code for prog assignment 2