book problem now passes all tests
refactored a bit to make it prettier