Improving the code to better perform