Refactoring of code for Assignment 2