使用 NumPy 数据类型的 Python 字典查找速度

Question

背景

我在 NumPy 数组中有很多数字消息代码，我需要快速将它们转换为字符串。我在性能方面遇到了一些问题，想了解原因以及如何快速解决。

一些基准

I - 简单的方法

import numpy as np

# dictionary to use as the lookup dictionary
lookupdict = {
     1: "val1",
     2: "val2",
    27: "val3",
    35: "val4",
    59: "val5" }

# some test data
arr = np.random.choice(lookupdict.keys(), 1000000)

# create a list of words looked up
res = [ lookupdict[k] for k in arr ]

字典查找占用了我咖啡休息时间的大部分时间，758 毫秒。 (我也试过 res = map(lookupdict.get, arr) 但那更糟。)

II - 没有 NumPy

import random

# dictionary to use as the lookup dictionary
lookupdict = {
     1: "val1",
     2: "val2",
    27: "val3",
    35: "val4",
    59: "val5" }

# some test data
arr = [ random.choice(lookupdict.keys()) for _ in range(1000000) ]

# create a list of words looked up
res = [ lookupdict[k] for k in arr ]

计时结果变化相当大，为 76 毫秒!

应该注意的是，我对查找的时间很感兴趣。随机生成只是为了创建一些测试数据。花不花时间没意思。此处给出的所有基准测试结果仅针对一百万次查找。

III - 将 NumPy 数组转换为列表

我的第一个猜测是这与列表与数组问题有关。但是，通过修改 NumPy 版本以使用列表:

res = [ lookupdict[k] for k in list(arr) ]

给了我 778 毫秒，其中大约 110 毫秒用于转换列表，570 毫秒用于查找。所以，查找要快一点，但总时间是一样的。

IV - 从np.int32到int的类型转换

由于唯一的其他区别似乎是数据类型(np.int32 与 int)，我尝试即时转换类型。这有点愚蠢，因为 dict 可能也是这样做的:

res = [ lookupdict[int(k)] for k in arr ]

然而，这似乎做了一些有趣的事情，因为时间下降到 266 毫秒。似乎几乎但不完全相同的数据类型在字典查找中玩弄了讨厌的把戏，而且 dict 代码在转换方面不是很有效。

V - 字典键转换为 np.int32

为了对此进行测试，我修改了 NumPy 版本以在字典键和查找中使用完全相同的数据类型:

import numpy as np

# dictionary to use as the lookup dictionary
lookupdict = {
     np.int32(1): "val1",
     np.int32(2): "val2",
    np.int32(27): "val3",
    np.int32(35): "val4",
    np.int32(59): "val5" }

# some test data
arr = np.random.choice(lookupdict.keys(), 1000000)

# create a list of words looked up
res = [ lookupdict[k] for k in arr ]

这提高到 177 毫秒。这不是微不足道的改进，而是与 76 毫秒相去甚远。

VI - 使用 int 对象的数组转换

import numpy as np

# dictionary to use as the lookup dictionary
lookupdict = {
     1: "val1",
     2: "val2",
    27: "val3",
    35: "val4",
    59: "val5" }

# some test data
arr = np.array([ random.choice(lookupdict.keys()) for _ in range(1000000) ], 
               dtype='object')

# create a list of words looked up
res = [ lookupdict[k] for k in arr ]

这给出了 86 毫秒，这已经非常接近原生 Python 的 76 毫秒。

结果总结

1. dict 键 int，使用 int 进行索引(原生 Python):76 毫秒
2. 字典键 int，使用 int 对象进行索引 (NumPy):86 毫秒
3. dict 键 np.int32，使用 np.int32 进行索引:177 毫秒
4. 字典键 int，使用 np.int32 进行索引:758 毫秒

问题(S)

为什么？我该怎么做才能尽可能快地进行字典查找？我的输入数据是一个 NumPy 数组，所以到目前为止最好的(最快但丑陋的)是将字典键转换为 np.int32。 (不幸的是，dict 键可能分布在很大范围的数字上，因此逐个数组索引不是一个可行的替代方案。虽然很快，10 毫秒。)

Answer 1

正如您所怀疑的，它是 int32.__hash__ 的错，x11 和 int.__hash__ 一样慢:

%timeit hash(5)
10000000 loops, best of 3: 39.2 ns per loop
%timeit hash(np.int32(5))
1000000 loops, best of 3: 444 ns per loop

(int32 类型是用 C 语言实现的。如果你真的很好奇，你可以深入研究源代码，找出它在那里做了什么，花了这么长时间)。

---

编辑:

第二个减慢速度的部分是字典查找中的隐式==比较:

a = np.int32(5)
b = np.int32(5)
%timeit a == b  # comparing two int32's
10000000 loops, best of 3: 61.9 ns per loop
%timeit a == 5  # comparing int32 against int -- much slower
100000 loops, best of 3: 2.62 us per loop

这解释了为什么您的 V 比 I 和 IV 快得多。当然，坚持使用全 int 解决方案会更快。

---

在我看来，您有两个选择:

1. 坚持使用纯 int 类型，或者在 dict-lookup 之前转换为 int
2. 如果最大代码值不是太大，和/或内存不是问题，您可以将字典查找换成列表索引，这不需要散列ing。

例如:

lookuplist = [None] * (max(lookupdict.keys()) + 1)
for k,v in lookupdict.items():
    lookuplist[k] = v

res = [ lookuplist[k] for k in arr ] # using list indexing

(编辑:您可能还想在这里试验 np.choose)