Refactor the code to improve performance