More (hopefully) code improvements for assignment 2