code tidyup / comments, and use defaultdict