python - 在数组中搜索一个值，然后从其他相同长度的数组/ndarrays 中打印相应的值

Question

我有一个低效的算法需要改进。

本质上我有数组:

ids = [134543, ..., 234]
a = [123, ..., 3546]
b = [[435, ..., 549][245, ..., 4986]]

所有相同的长度(即 ids.shape = (600000,) a.shape = (600000,) b.shape = (600000, 2))

和一个更小的 id 数组(长度大约 100)

ids_important = [345, ..., 549]

我想在我的ids数组中找到重要id的索引，然后输出a和b中对应的元素。

我目前的算法是:

for i in range(ids_important.shape[0]):
    for j in range(ids.shape[0]):
        if(ids[j] == ids_important[i]):
           print(a[j])
           print(b[j,0])

考虑到数组的大小，这个算法非常慢。有人告诉我可以使用屏蔽数组改进它，但一直无法弄清楚如何实现它。非常感谢您的帮助。

Answer 1

这可能不是最节省空间的方法(分配一个 len(important_ids) x len(ids) 数组)但它应该比 多快另一个答案和您原来的方法，因为它利用了 numpy 向量化(并避免了慢循环):

import numpy as np
ids = np.random.randint(0, 15, size=20)
important_ids = np.random.randint(5, 9, size=10)
a = np.random.randint(10, 30, size=50)
b = np.random.randint(10, 30, size=(50, 2))

equals = np.argwhere(ids[None,:] == important_ids[:, None])
inds = equals[:, 1]
print(a[inds])
print(b[tuple(inds), 0])