solve problem 4 + small refactoring