Update with a slightly better solution.
Refactor the code and update with comments.