Search Problem #196, improving code