Fix while loop in problem 7 in Python