Optimised solution to problem 4
Solved problem 5