Adding some logic to improve speed