Assignment completion

Added code to complete the assignment