如何在有序数组中超越numpy的in1d掩码函数？

Question

这个操作需要尽可能快地完成，因为实际的数组包含数百万个元素。这是一个简单版本的问题。

我有一个随机的、唯一的整数数组（通常有数百万个元素）。

totalIDs = [5,4,3,1,2,9,7,6,8 ...]

我还有另一个数组（通常有几万个）也是唯一的整数，我可以用它来创建一个掩码。

subsampleIDs1 = [5,1,9]
subsampleIDs2 = [3,7,8]
subsampleIDs3 = [2,6,9]
...

我可以使用numpy来完成这个操作：

mask = np.in1d(totalIDs,subsampleIDs,assume_unique=True)

然后，我可以利用这个掩码从另一个数组中提取我想要的信息（比如说，第一列是我想要的）。

variable = allvariables[mask][:,0]

现在，考虑到两个数组中的ID都是唯一的，有没有办法显著加快这个过程？为几千个点（subsampleIDs）在数百万个ID（totalIDs）中构建掩码需要很长时间。

我想过先遍历一次，然后写出一个索引的二进制文件（以加快未来的搜索）。

for i in range(0,3):
    mask = np.in1d(totalIDs,subsampleIDs,assume_unique=True)
    index[mask] = i

其中X在subsampleIDsX中。然后我可以直接这样做：

for i in range(0,3):
    if index[i] == i:
        rowmatch = i
        break

variable = allvariables[rowmatch:len(subsampleIDs),0]

对吧？但这也很慢，因为循环中有一个条件判断来查找第一次匹配的情况。有没有更快的方法来找到一个数字在有序数组中第一次出现的位置，这样条件判断就不会拖慢循环速度？

Answer 1

通常，这种索引最好是用数据库来做，并且要正确设置列的索引。

另一个想法是先对 totalIDs 进行一次排序，作为预处理步骤，然后自己实现一个版本的 in1d，这样就可以避免每次都排序。numpy 中的 in1d 实现（至少在我安装的版本中）相对简单，应该很容易复制和修改。

编辑：

或者，更好的方法是使用桶排序（或基数排序）。这样可以让你的时间复杂度达到 O(N+M)，其中 N 是 totalIDs 的大小，M 是 sampleIDs 的大小（你可以通过改变桶的数量来调整一个常数）。在这里，你也只需要对 totalIDs 进行一次桶的划分，这样可以得到一个很不错的 O(N+M1+M2+...)。

不幸的是，我不知道 numpy 有没有这样的实现，但我找到了一些资料：http://en.wikipedia.org/wiki/Radix_sort#Example_in_Python

Answer 2

我建议你使用Pandas中的DataFrame。DataFrame的索引是totalIDs，你可以通过 df.ix[subsampleIDs] 来选择subsampleIDs。

首先创建一些测试数据：

import numpy as np
N = 2000000
M = 5000
totalIDs = np.random.randint(0, 10000000, N)
totalIDs = np.unique(totalIDs)
np.random.shuffle(totalIDs)
v1 = np.random.rand(len(totalIDs))
v2 = np.random.rand(len(totalIDs))

subsampleIDs = np.random.choice(totalIDs, M)
subsampleIDs = np.unique(subsampleIDs)
np.random.shuffle(subsampleIDs)

然后把你的数据转换成DataFrame：

import pandas as pd
df = pd.DataFrame(data = {"v1":v1, "v2":v2}, index=totalIDs) 
df.ix[subsampleIDs]

DataFrame使用哈希表来将索引和它的位置对应起来，这样速度非常快。