UPDATE to 'solved pset 3'