Implemented correct solution for problem 5