比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)

然后我可以使用掩码提取我想要的另一个数组的信息(比如第0列包含我想要的那个).

variable = allvariables [mask] [:,0]

既然ID在两个阵列中都是唯一的,那么有没有办法显着提高速度.构建与数百万个ID(totalIDs)匹配的几千个点(subsampleIDs)的掩码需要很长时间.

我曾经想过要经历一次并写出一个索引的二进制文件(以加快未来的搜索).

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

我建议你在 Pandas 中使用 DataFrame。DataFrame 的索引是totalIDs，您可以通过以下方式选择subsampleIDs：df.ix[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 使用哈希表将索引映射到它的位置，速度非常快。