Minor addition to problem6.py