python - 如何匹配两个数组的元素并返回一个值

Question

我正在寻求帮助来改进我的代码。我有两个不同大小的数组，我将一个数组的值分配给另一个数组的索引。例如

a = np.array([[1,5],[2,8],[3,2],[4,6]])
b = np.array([[1],[1],[2],[4],[4],[2],[3],[2],[3]])

我正在计算第一列中的项目，然后将 a 的第二列中的值分配给 b 中的相应行。

c = np.zeros([len(b),1])
for i in tqdm(range(len(b))):
    for j in range(len(a)):
        if b[i,0]==a[j,0]:
            c[i] = a[j,1]

返回

c = [[5],[5],[8],[6],[6],[8],[2],[8],[2]]

问题是我有一个非常大的数据集，for 循环需要很长时间才能运行。任何建议将不胜感激。谢谢。

Answer 1

a 是否总是格式 [[1, ...], [2, ...], [3, ...], ...]?

如果是，那么您可以通过不遍历 a 而只是索引它来节省时间。例如:

a = np.array([[1,5],[2,8],[3,2],[4,6]])
b = np.array([[1],[1],[2],[4],[4],[2],[3],[2],[3]])

c = np.array([[a[i[0] - 1][1]] for i in b])
# c = [[5], [5], [8], [6], [6], [8], [2], [8], [2]]

这将采用 b 时间大小的顺序，而不是 a 时间大小乘以 b 时间大小的顺序。