Solution for problem 3 improved