optimized the code a bit