Small refactoring to Problem 8