Python字典查找速度与NumPy数据类型

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），我尝试在运行时转换类型。这有点傻，因为字典可能也会这样做：

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

不过，这似乎产生了一些有趣的效果，因为时间降到了266毫秒。看起来几乎相同的数据类型在字典查找时会造成麻烦，而字典代码在转换时效率不高。

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. 字典键为 int，用 int 索引（原生Python）：76毫秒
2. 字典键为 int，用 int 对象索引（NumPy）：86毫秒
3. 字典键为 np.int32，用 np.int32 索引：177毫秒
4. 字典键为 int，用 np.int32 索引：758毫秒

问题

为什么会这样？我该怎么做才能让字典查找尽可能快？我的输入数据是一个NumPy数组，所以到目前为止，最快（但不太优雅）的方法是把字典键转换成 np.int32。（不幸的是，字典键可能分布在很大的数字范围内，所以逐个数组索引并不是一个可行的替代方案。虽然这样会很快，10毫秒。）

Answer 1

当你的键是整数时，这似乎是最快的方法。只需要使用数组索引就可以了。

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)
lookup_array=np.array([None if i not in lookupdict else lookupdict[i] for i in range(60)])
%timeit lookup_array[arr]

每次循环大约需要 5.3 毫秒，误差范围是 ± 73.7 微秒（这是7次运行的平均值和标准差，每次运行100次）

Answer 2

这里有一个使用Pandas的解决方案，速度提高了五倍：

import numpy as np
import pandas as pd

# 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
%timeit res = [ lookupdict[k] for k in arr ]
%timeit res_pd = pd.Series(lookupdict).reindex(arr).values
print all(res == res_pd)

10 loops, best of 3: 192 ms per loop
10 loops, best of 3: 35.3 ms per loop
True

这意味着每个元素的平均处理时间是35纳秒，所以在原生Python中很难再快了。如果你对Pandas不太了解，可以把Series对象想象成一个有序字典或者一个带索引的数组，它可以从一个普通的Python字典构建而来。reindex这个方法提供了非常快速的查找；我不太清楚它是怎么做到的，因为我不是很有经验的程序员，但它可能是用C或Cython写的。也许你可以查看一下源代码，找出一个更快的解决方案来解决你的问题。最后，values属性只是返回Series背后的数组。

编辑：
这里有一个纯粹使用numpy的解决方案，效果几乎和Pandas一样好：

keys = np.array(lookupdict.keys())
strings = np.array(lookupdict.values())
%timeit res_np = strings[(np.atleast_2d(arr).T == keys).argmax(axis=1)]
10 loops, best of 3: 44.6 ms per loop

print all(res == res_np)
True

Answer 3

这很有意思，我可能找到了我问题的答案。

另外一种方法是把数组转换成列表。看起来如果方法对了，这样做效果非常好。比如：

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

这个操作耗时778毫秒。

但是这个：

res = [ lookupdict[k] for k in arr.tolist() ]

耗时只有86毫秒。

这里的技术原因是，arr.tolist会把数组转换成int对象，而list(arr)则是创建一个np.int32对象的列表。

Answer 4

正如你所猜测的，问题出在 int32.__hash__ 上，它的速度和 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 类型，或者在查字典之前先转换为 int
2. 如果最大的代码值不是太大，或者内存不是问题，你可以用列表索引来替代字典查找，这样就不需要 hash 了。

例如：

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）

Answer 5

根据我的测试，你的 II - Without NumPy 的速度比 I 慢很多。

In [11]: timeit [lookupdict[k] for k in np.random.choice(lookupdict.keys(),1000000)]
1 loops, best of 3: 658 ms per loop

In [12]: timeit [lookupdict[k] for k in [np.random.choice(lookupdict.keys()) for _ in range(1000000)]]
1 loops, best of 3: 8.04 s per loop

但是如果你跳过查找，直接对值进行 choice 操作，你会节省更多时间。

In [34]: timeit np.random.choice(lookupdict.values(),1000000)
10 loops, best of 3: 85.3 ms per loop

---

好吧，我们来专注于查找的部分：

In [26]: arr =np.random.choice(lookupdict.keys(),1000000)

In [27]: arrlist=arr.tolist()

In [28]: timeit res = [lookupdict[k] for k in arr]
1 loops, best of 3: 583 ms per loop

In [29]: timeit res = [lookupdict[k] for k in arrlist]
10 loops, best of 3: 120 ms per loop

In [30]: timeit res = [lookupdict[k] for k in list(arr)]
1 loops, best of 3: 675 ms per loop

In [31]: timeit res = [lookupdict[k] for k in arr.tolist()]
10 loops, best of 3: 156 ms per loop

In [32]: timeit res = [k for k in arr]
1 loops, best of 3: 215 ms per loop

In [33]: timeit res = [k for k in arrlist]
10 loops, best of 3: 51.4 ms per loop

In [42]: timeit arr.tolist()
10 loops, best of 3: 33.6 ms per loop

In [43]: timeit list(arr)
1 loops, best of 3: 264 ms per loop

首先观察到，遍历一个 np.array 的速度比遍历一个相同的列表要慢。

其次，使用 list(arr) 的速度比 arr.tolist() 慢。list() 似乎有两个问题：它本身就比较慢，而且里面的元素是 np.int32 类型。