如何使用这个用 Numpy 构建的集合字典提高性能？

如何解决如何使用这个用 Numpy 构建的集合字典提高性能？

假设我们有 1000 万只袜子，每只袜子的颜色为 sockcolor[i] 并存储在 drawer[i] 中。我想数一数 20 个抽屉中的每个抽屉有多少不同种颜色。

为此，我使用了一个包含集合的字典（集合在引擎盖下使用了一个哈希表，它可以方便地计算唯一元素）。

以下代码有效，但速度很慢（1000 万只袜子约 10 秒）。

我们如何使用 Numpy 技术（矢量化？）或避免 for 循环来加速此计算？

import numpy as np
N = 10*1000*1000
drawer = np.random.randint(0,20,N)            # 20 drawers
sockcolor = np.random.randint(0,100*1000,N)   # 100k different colors
d = {}

for i,k in enumerate(drawer):
    if k not in d:  # can be simplified with collections.defaultdict but no gain here
        d[k] = set()
    d[k].add(sockcolor[i])

for k,s in d.items():
    print(k,len(s))

输出：

解决方法

您基本上已经有了从抽屉到袜子颜色的映射，但它们是随机的，您想按抽屉编号组织它们。

最简单的方法是先按抽屉编号对它们进行排序：

drawer_sort = np.argsort(drawer)
drawer = drawer[drawer_sort]
sockcolor = sockcolor[drawer_sort]

既然它们已排序，则无需查找抽屉编号重复项，您只需找到抽屉编号更改的索引，以形成范围，即这些：

changes,= np.where(drawer[1:]-drawer[:-1])
starts = np.concatenate([[0],changes+1])
ends = np.concatenate([changes,[len(drawer)]])

现在你可以创建你的字典了：

result = {drawer[start]: sockcolor[start:end] for start,end in zip(starts,ends)}

这样，在 Python 中唯一完成的迭代是最后一行，如果有少量不同的 drawer 值（在您的情况下不超过 20），这将非常快。

结果仍然可能有重复的 sockcolor 值，但这在 numpy 中很容易解决：

result = {drawer: np.unique(sockcolors) for drawer,sockcolors in result.items()}

您的缓慢是因为没有使用序列的内置功能。仅迭代单个袜子一次。相反，将袜子颜色（不是单个袜子的索引）分配给抽屉。然后从每个抽屉的内容制作一组：一次批发操作，而不是增量 set.add，这对于您的目的来说相对较慢。