Cleaned up code to follow Hackerrank's input format