updated problem 2 solution