problem set 4 code added