improving my code