Solution (almost) to problem 2