Code modified

Modified code for better performance