Add code for problem #4 and basic unit tests