Improved code to pass mutant.