Improved code, removed unnecessary variables