Slightly improved the answer for problem 2