Updated code

Updated code with feedback from community