Working Solution to Problem 2

Solution to problem 2, a bit more concise than Problem 1