nShift interview question

- improve the code by solving some redundant implementation - fix unit tests