python - 禁止在 collections.defaultdict 中添加键

Question

当在 defaultdict 对象中查询丢失的键时，该键自动添加到字典中:

from collections import defaultdict

d = defaultdict(int)
res = d[5]

print(d)
# defaultdict(<class 'int'>, {5: 0})
# we want this dictionary to remain empty

但是，通常我们只想在显式或隐式分配键时才添加键:

d[8] = 1  # we want this key added
d[3] += 1 # we want this key added

一个用例是简单计数，以避免 collections.Counter 的更高开销，但通常也可能需要此功能。

---

反例 [双关语]

这是我想要的功能:

from collections import Counter
c = Counter()
res = c[5]  # 0
print(c)  # Counter()

c[8] = 1  # key added successfully
c[3] += 1 # key added successfully

但是 Counter 明显比 defaultdict(int) 慢。我发现性能下降通常比 defaultdict(int) 慢 2 倍。

此外，显然Counter只能与defaultdict中的int参数相比，而defaultdict可以带list、set等

---

有没有办法有效地实现上述行为；例如，通过继承 defaultdict?

---

基准测试示例

%timeit DwD(lst)           # 72 ms
%timeit dd(lst)            # 44 ms
%timeit counter_func(lst)  # 98 ms
%timeit af(lst)            # 72 ms

测试代码:

import numpy as np
from collections import defaultdict, Counter, UserDict

class DefaultDict(defaultdict):
    def get_and_forget(self, key):
        _sentinel = object()
        value = self.get(key, _sentinel)

        if value is _sentinel:
            return self.default_factory()
        return value

class DictWithDefaults(dict):
    __slots__ = ['_factory']  # avoid using extra memory

    def __init__(self, factory, *args, **kwargs):
        self._factory = factory
        super().__init__(*args, **kwargs)

    def __missing__(self, key):
        return self._factory()

lst = np.random.randint(0, 10, 100000)

def DwD(lst):
    d = DictWithDefaults(int)
    for i in lst:
        d[i] += 1
    return d

def dd(lst):
    d = defaultdict(int)
    for i in lst:
        d[i] += 1
    return d

def counter_func(lst):
    d = Counter()
    for i in lst:
        d[i] += 1
    return d

def af(lst):
    d = DefaultDict(int)
    for i in lst:
        d[i] += 1
    return d

关于赏金评论的注意事项:

@Aran-Fey's solution自提供赏金后已更新，因此请忽略赏金评论。

Answer 1

而不是乱搞 collections.defaultdict让它做我们想做的事，似乎更容易实现我们自己的:

class DefaultDict(dict):
    def __init__(self, default_factory, **kwargs):
        super().__init__(**kwargs)

        self.default_factory = default_factory

    def __getitem__(self, key):
        try:
            return super().__getitem__(key)
        except KeyError:
            return self.default_factory()

这按照你想要的方式工作:

d = DefaultDict(int)

res = d[5]
d[8] = 1 
d[3] += 1

print(d)  # {8: 1, 3: 1}

但是，对于可变类型，它可能会出现意外行为:

d = DefaultDict(list)
d[5].append('foobar')

print(d)  # output: {}

这可能是 defaultdict 在访问不存在的键时记住该值的原因。

---

另一种选择是扩展 defaultdict 并添加一个新方法来查找一个值而不用记住它:

from collections import defaultdict

class DefaultDict(defaultdict):
    def get_and_forget(self, key):
        return self.get(key, self.default_factory())

请注意，get_and_forget 方法每次都会调用 default_factory()，无论该键是否已存在于字典中。如果这是不可取的，您可以使用标记值来实现它:

class DefaultDict(defaultdict):
    def get_and_forget(self, key):
        _sentinel = object()
        value = self.get(key, _sentinel)

        if value is _sentinel:
            return self.default_factory()
        return value

这对可变类型有更好的支持，因为它允许您选择是否应该将值添加到字典中。