Updated solution to homework