* Add doctest and remove input() usage * Apply suggestions from code review Co-authored-by: Dhruv Manilawala <dhruvmanila@gmail.com>