禁止在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).我发现性能打击通常比慢2倍defaultdict(int).

此外,很明显Counter是只可比int的说法defaultdict,而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的解决方案自Bounty发布以来已经更新,所以请忽略Bounty的评论.

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())

请注意,无论密钥是否已存在于dict中,该get_and_forget方法default_factory()都会每次调用.如果这是不合需要的,您可以使用sentinel值来实现它:

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

这样可以更好地支持可变类型,因为它允许您选择是否应该将值添加到dict中.

Answer 2

如果只想dict在访问不存在的键时返回默认值的a ，则可以简单地子类化dict并实现__missing__：

object.__missing__(self, key)

当不在字典中时，由调用dict.__getitem__()以实现self[key]dict子类key。

看起来像这样：

class DictWithDefaults(dict):
    # not necessary, just a memory optimization
    __slots__ = ['_factory']  

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

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

在这种情况下，我使用了defaultdict-like方法，因此您必须传入，factory在调用时应提供默认值：

>>> dwd = DictWithDefaults(int)
>>> dwd[0]  # key does not exist
0 
>>> dwd     # key still doesn't exist
{}
>>> dwd[0] = 10
>>> dwd
{0: 10}

当您进行分配（显式或隐式）时，该值将添加到字典中：

>>> dwd = DictWithDefaults(int)
>>> dwd[0] += 1
>>> dwd
{0: 1}

>>> dwd = DictWithDefaults(list)
>>> dwd[0] += [1]
>>> dwd
{0: [1]}

您想知道它collections.Counter是如何做到的，从CPython 3.6.5开始，它还使用了__missing__：

class Counter(dict):
    ...

    def __missing__(self, key):
        'The count of elements not in the Counter is zero.'
        # Needed so that self[missing_item] does not raise KeyError
        return 0

    ...

更好的性能？！

您提到速度是值得关注的，因此您可以使该类成为C扩展类（假设您使用CPython），例如使用Cython（我正在使用Jupyter magic命令创建扩展类）：

%load_ext cython

%%cython

cdef class DictWithDefaultsCython(dict):
    cdef object _factory

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

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

基准测试

根据您的基准：

from collections import Counter, defaultdict

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

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

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

考虑到这只是计数，如果不仅仅使用Counter初始化程序就不包含基准，将是（不可原谅的）疏忽。

我最近写了一个小型基准测试工具，我认为它可能会派上用场（但您也可以使用%timeit）：

from simple_benchmark import benchmark
import random

sizes = [2**i for i in range(2, 20)]
unique_lists = {i: list(range(i)) for i in sizes}
identical_lists = {i: [0]*i for i in sizes}
mixed = {i: [random.randint(0, i // 2) for j in range(i)]  for i in sizes}

functions = [d_py, d_cy, d_dd, d_c, Counter]

b_unique = benchmark(functions, unique_lists, 'list size')
b_identical = benchmark(functions, identical_lists, 'list size')
b_mixed = benchmark(functions, mixed, 'list size')

结果如下：

import matplotlib.pyplot as plt

f, (ax1, ax2, ax3) = plt.subplots(1, 3, sharey=True)
ax1.set_title('unique elements')
ax2.set_title('identical elements')
ax3.set_title('mixed elements')
b_unique.plot(ax=ax1)
b_identical.plot(ax=ax2)
b_mixed.plot(ax=ax3)

请注意，它使用对数对数刻度来更好地查看差异：

长期以来，迭代Counter(iterable)速度最快。DictWithDefaultCython并defaultdict等于（DictWithDefault大多数情况下速度稍快，即使在这里看不见）DictWithDefault，然后Counter是手动for-loop。有趣的Counter是最快和最慢。

隐式添加返回的值，如果它是modifie

我忽略的一个事实是，它defaultdict与可变类型有很大的不同，这是因为期望的“只是返回默认值而不保存它”：

>>> from collections import defaultdict
>>> dd = defaultdict(list)
>>> dd[0].append(10)
>>> dd
defaultdict(list, {0: [10]})

>>> dwd = DictWithDefaults(list)
>>> dwd[0].append(10)
>>> dwd
{}

这意味着当您希望修改后的值在字典中可见时，实际上需要设置元素。

但是，这有点让我着迷，所以我想分享一种您如何进行这项工作的方式（如果需要）。但这只是一个快速测试，仅适用于append使用代理的呼叫。请不要在生产代码中使用它（从我的角度来看，这只是具有娱乐价值）：

from wrapt import ObjectProxy

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

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

    def __missing__(self, key):
        ret = self._factory()
        dict_ = self

        class AppendTrigger(ObjectProxy):
            def append(self, val):
                self.__wrapped__.append(val)
                dict_[key] = ret

        return AppendTrigger(ret)

那是一个字典，它返回一个代理对象（而不是真正的默认值），并且它重载了一个方法，如果调用该方法，则将返回值添加到字典中。它“有效”：

>>> d = DictWithDefaultsFunky(list)
>>> a = d[10]
>>> d
[]

>>> a.append(1)
>>> d
{10: [1]}

但这确实有一些陷阱（可以解决，但这只是概念验证，因此我在这里不再尝试）：

>>> d = DictWithDefaultsFunky(list)
>>> a = d[10]
>>> b = d[10]
>>> d
{}
>>> a.append(1)
>>> d
{10: [1]}
>>> b.append(10)
>>> d  # oups, that overwrote the previous stored value ...
{10: [10]}

如果您确实想要类似的东西，则可能需要实现一个类，该类真正跟踪值内的更改（而不仅仅是append调用）。

如果要避免隐式分配

如果您不喜欢+=或类似的操作将值添加到字典中的事实（与上一个示例甚至试图以非常隐式的方式添加值相反），那么您可能应该将其实现为方法而不是特殊的方法方法。

例如：

class SpecialDict(dict):
    __slots__ = ['_factory']

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

    def get_or_default_from_factory(self, key):
        try:
            return self[key]
        except KeyError:
            return self._factory()

>>> sd = SpecialDict(int)
>>> sd.get_or_default_from_factory(0)  
0
>>> sd  
{}
>>> sd[0] = sd.get_or_default_from_factory(0) + 1
>>> sd  
{0: 1}

这与Aran-Feys答案的行为类似，但是get它使用tryand catch方法代替了哨兵。