fixed small bug in problem 6