我们有data 和reference 作为-
In [375]: data
Out[375]: array([30, 20, 30, 10, 20, 10, 20, 10, 30, 20, 20, 30, 30, 10, 30])
In [376]: reference
Out[376]: array([20, 10, 30])
让我们考虑一下reference的排序版本-
In [373]: np.sort(reference)
Out[373]: array([10, 20, 30])
现在,我们可以使用np.searchsorted 来找出每个data 元素在这个排序版本中的位置,就像这样 -
In [378]: np.searchsorted(np.sort(reference), data, side='left')
Out[378]: array([2, 1, 2, 0, 1, 0, 1, 0, 2, 1, 1, 2, 2, 0, 2], dtype=int64)
如果我们运行原始代码,预期的输出结果是 -
In [379]: indexes
Out[379]: array([2, 0, 2, 1, 0, 1, 0, 1, 2, 0, 0, 2, 2, 1, 2])
可以看出,searchsorted 输出很好,除了其中的0's 必须是1s 和1's 必须更改为0's。现在,我们已经开始计算,reference 的排序版本。因此,要进行0's 到1's 的更改,反之亦然,我们需要引入用于排序reference 的索引,即np.argsort(reference)。这基本上就是矢量化无循环或无字典的方法!所以,最终的实现看起来像这样 -
# Get sorting indices for reference
sort_idx = np.argsort(reference)
# Sort reference and get searchsorted indices for data in reference
pos = np.searchsorted(reference[sort_idx], data, side='left')
# Change pos indices based on sorted indices for reference
out = np.argsort(reference)[pos]
运行时测试 -
In [396]: data = np.random.randint(0,30000,150000)
...: reference = np.unique(data)
...: reference = reference[np.random.permutation(reference.size)]
...:
...:
...: def org_approach(data,reference):
...: indexes = np.zeros_like(data, dtype=int)
...: for i in range(data.size):
...: indexes[i] = np.where(data[i] == reference)[0]
...: return indexes
...:
...: def vect_approach(data,reference):
...: sort_idx = np.argsort(reference)
...: pos = np.searchsorted(reference[sort_idx], data, side='left')
...: return sort_idx[pos]
...:
In [397]: %timeit org_approach(data,reference)
1 loops, best of 3: 9.86 s per loop
In [398]: %timeit vect_approach(data,reference)
10 loops, best of 3: 32.4 ms per loop
验证结果 -
In [399]: np.array_equal(org_approach(data,reference),vect_approach(data,reference))
Out[399]: True