Updated code for efficiency