Edited code for efficiency