引数がある初期化関数を用いた, defaultdictの実装 (Python)

目的

keyを引数とする関数で初期化を行うようなdefaultdictの実装

通常のdefaultdict

通常のdefaultdict. dictの初期化が, 与えた関数に従って行われます.

>>> from collections import defaultdict
>>> dd = defaultdict(int)
>>> dd[1] += 1  # dd[1] = (0 + 1)
>>> dd
defaultdict(<class 'int'>, {1: 1})

lambda式関数などで, 自作の関数を用いることもできます.

>>> dd = defaultdict(lambda: 10)
>>> dd[1] += 1  # dd[1] = (10 + 1)
>>> dd
defaultdict(<class 'int'>, {1: 11})

しかし, defaultdictの関数に用いることができるのは, 引数なしの関数のみとなっており,
例えば, $f(x) = 2x$ という関数で初期化を行いたい場合, 通常のdefaultdictで実現することはできません(たぶん).

引数ありのdefaultdictの実装

まず, 通常のdefaultdictを簡単に自作しようとすると, 以下のようになります.

# defualt dictの実装
class SimpleDefaultdict(dict):
    def __init__(self, func):
        self.func = func
    def __getitem__(self, item):
        if item not in self:
            self[item] = self.func()
        return super().__getitem__(item)

dictを継承して実装していて, __getitem__ で[]参照されたときの処理を定義しています. itemが辞書dになければ, 与えた関数の返り値で初期化を行う具合ですね.
余談ですが, もし __getitem__の返り値をreturn self[item]にしてしまうと, self[item]の計算時に__getitem__が呼ばれ, またその計算時に__getitem__が呼ばれて...が繰り返されて無限再帰となってしまうので注意が必要です.

これを踏まえると, 引数ありの関数でのdefualt dictは次のようにすれば作成できます.

# 引数ありの場合
class MyDefaultdict(dict):
    def __init__(self, func):
        self.func = func
    def __getitem__(self, item):
        if item not in self:
            self[item] = self.func(item)  # <-- ここが変わった
        return super().__getitem__(item)

初期化の際に関数に引数を与えています. 実際に$f(x) = 2x$ という関数で初期化を行う場合

>>> mdd = MyDefaultdict(lambda x: 2*x)
>>> mdd[1]
2
>>> mdd
{1: 2}

できていますね.

まとめ

dictクラスを継承して, 引数ありの関数によるdefaultdictを実装しました.