匹配numpy数组中的元素

Question

我有两个 numpy 数组。第一个 Z1 大约有 300,000 行长和 3 列宽。第二个，Z2，大约有 200,000 行和 300 列。每个 Z1 和 Z2 的每一行都有一个标识号（10 位）。 Z2 包含 Z1 中的项目的子集，我想根据 10 位标识号将 Z2 中的行与 Z1 中的伙伴匹配，然后从 Z1 中取出第 2 列和第 3 列并将它们插入到 Z2 的末尾相应的行。

Z1 和 Z2 都没有任何特定的顺序。

我想出的唯一方法是遍历数组，这需要几个小时。有没有更好的方法在 Python 中做到这一点？

谢谢！

Answer 1

从您的问题中我了解到，10 位标识符存储在第 1 列中，对吗？

这不是很容易理解，很多间接发生，但最后unsorted_insert 有Z1 中Z2 的每个标识符的行号是

sort_idx = np.argsort(Z1[:, 0])
sorted_insert = np.searchsorted(Z1[:, 0], Z2[:, 0], sorter=sort_idx)
# The following is equivalent to unsorted_insert = sort_idx[sorted_insert] but faster
unsorted_insert = np.take(sort_idx, sorted_insert)

所以现在我们需要做的就是获取这些行的最后两列并将它们堆叠到Z2 数组中：

new_Z2 = np.hstack((Z2, Z1[unsorted_insert, 1:]))

---

一个运行没有问题的虚构示例：

import numpy as np

z1_rows, z1_cols = 300000, 3
z2_rows, z2_cols = 200000, 300

z1 = np.arange(z1_rows*z1_cols).reshape(z1_rows, z1_cols)

z2 = np.random.randint(10000, size=(z2_rows, z2_cols))
z2[:, 0] = z1[np.random.randint(z1_rows, size=(z2_rows,)), 0]

sort_idx = np.argsort(z1[:, 0])
sorted_insert = np.searchsorted(z1[:, 0], z2[:, 0], sorter=sort_idx)
# The following is equivalent to unsorted_insert = sort_idx[sorted_insert] but faster
unsorted_insert = np.take(sort_idx, sorted_insert)
new_z2 = np.hstack((z2, z1[unsorted_insert, 1:]))

还没有计时，但整个事情似乎在几秒钟内完成。