Made changes according to spec written in the challenge