Refactor of problem 3