Update answer2.py

This is the answer to question 2, i think it still has a bug.