Search code examples
pythondefaultdict

More Pythonic way of counting things in a heavily nested defaultdict


My code currently has to count things in a heavily nested dict into another. I have items that need to be indexed by 3 values and then counted. So, before my loop, I initialize a nested defaultdict like so:

from collections import defaultdict

type_to_count_dic = defaultdict(
        lambda: defaultdict(
            lambda: defaultdict(int)
        )
    )

Which allows me to count the items within a tight loop like so:

for a in ...:
    for b in ...:
        for c in ...:
            type_to_count_dic[a][b][c] += 1

I feel like initializing all those defaultdicts feels a lot like making a type declaration in something like Java. Is there a more idiomatic/Pythonic way of doing something like this?


Solution

  • from collections import defaultdict
    
    class _defaultdict(defaultdict):
        def __add__(self, other):
            return other
    
    def CountTree():
        return _defaultdict(CountTree)
    

    >>> t = CountTree()
    >>> t['a']
    defaultdict(<function CountTree at 0x9e5c3ac>, {})
    >>> t['a']['b']['c'] += 1
    >>> print t['a']['b']['c']
    1