Reduced amount of code and made it handle N floors