Slightly improved logic in problem 17