Modify setup.py

I think use `defaultdict` is simpler