python - 快速 python 算法(在 numpy 或 pandas 中？)查找与另一个数组中的元素匹配的数组元素的索引

Question

我正在寻找一种快速方法来确定两个数组的交叉匹配索引，定义如下。

我有两个非常大的(>1e7 个元素)结构化数组，一个叫做members，另一个叫做groups。两个数组都有一个 groupID 列。 groups 数组的 groupID 条目是唯一的，members 数组的 groupID 条目不是。

groups 数组有一个名为mass 的列。 members 数组有一个名为 groupmass 的列(当前为空)。我想将正确的 groupmass 分配给 members 的那些元素，其中 groupID 与其中一个组匹配。这将通过以下方式完成:

members['groupmass'][idx_matched_members] = groups['mass'][idx_matched_groups]

所以我需要一个快速例程来计算两个索引数组 idx_matched_members 和 idx_matched_groups。这种任务似乎很常见，以至于像 numpy 或 pandas 这样的包似乎很可能会有一个优化的解决方案。有谁知道专业开发的、自制的或其他方式的解决方案？

Answer 1

这可以通过 pandas 使用 map 将一列的数据映射到另一列的数据来完成。这是一个示例数据示例:

members = pandas.DataFrame({
    'id': np.arange(10),
    'groupID': np.arange(10) % 3,
    'groupmass': np.zeros(10)
})

groups = pandas.DataFrame({
    'groupID': np.arange(3),
    'mass': np.random.randint(1, 10, 3)
})

这为您提供了以下数据:

>>> members
   groupID  groupmass  id
0        0          0   0
1        1          0   1
2        2          0   2
3        0          0   3
4        1          0   4
5        2          0   5
6        0          0   6
7        1          0   7
8        2          0   8
9        0          0   9
>>> groups
   groupID  mass
0        0     3
1        1     7
2        2     4

然后:

>>> members['groupmass'] = members.groupID.map(groups.set_index('groupID').mass)
>>> members
   groupID  groupmass  id
0        0          3   0
1        1          7   1
2        2          4   2
3        0          3   3
4        1          7   4
5        2          4   5
6        0          3   6
7        1          7   7
8        2          4   8
9        0          3   9

如果你经常想使用 groupID 作为 groups 的索引，你可以永久设置它，这样你就不必每次都使用 set_index你这样做。