Question

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

我有两个非常大的（＆gt; 1e7元素）结构化数组，一个名为成员，另一个名为 groups 。两个阵列都有一个 groupID 列。 groups 数组的 groupID 条目是唯一的，成员数组的 groupID 条目不是唯一的。

groups 数组有一个名为 mass 的列。成员数组有一个名为 groupmass 的（当前为空）列。我想将 groupmass 分配给成员的那些元素，其中 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。

Answer 2

以下是仅使用mass设置numpy的示例。它确实使用迭代，因此对于大型数组，它不会很快。

仅10行，这比pandas等效快得多。但随着数据集变大（例如，M = 10000），pandas要好得多。 pandas的设置时间较长，但每行迭代时间要低得多。

生成测试数组：

dt_members = np.dtype({'names':['groupID','groupmass'], 'formats': [int, float]})
dt_groups =  np.dtype({'names':['groupID', 'mass'], 'formats': [int, float]})

N, M = 5, 10
members = np.zeros((M,), dtype=dt_members)    
groups = np.zeros((N,), dtype=dt_groups)
members['groupID'] = np.random.randint(101, 101+N, M)
groups['groupID'] = np.arange(101, 101+N)
groups['mass']  = np.arange(1,N+1)

def getgroup(id):
    idx = id==groups['groupID']
    return groups[idx]

members['groupmass'][:] = [getgroup(id)['mass'] for id in members['groupID']]

在python2中，迭代可以使用map：

members['groupmass'] =  map(lambda x: getgroup(x)['mass'], members['groupID'])

我可以通过最小化重复下标来提高速度约2倍，例如

def setmass(members, groups):
    gmass = groups['mass']
    gid = groups['groupID']
    mass = [gmass[id==gid] for id in members['groupID']]
    members['groupmass'][:] = mass

但如果groups['groupID']可以映射到arange(N)，那么我们可以大幅提升速度。通过将相同的映射应用于members['groupID']，它将成为一个简单的数组索引问题。

在我的示例数组中，groups['groupID']只是arange(N)+101。因此，映射只会减去最小值。

def setmass1(members, groups):
    members['groupmass'][:] = groups['mass'][members['groupID']-groups['groupID'].min()]

这比我之前的代码快300倍，比pandas解决方案（10000,500阵列）好8倍。

我怀疑pandas做了类似的事情。 pgroups.set_index('groupID').mass系列是mass，添加了.index属性。（我可以使用更通用的数组来测试它）

在更一般的情况下，对groups进行排序可能有所帮助，如有必要，可填写一些索引空白。

这是一个'矢量化'解决方案 - 没有迭代。但它必须计算一个非常大的矩阵（成员长度的组长度），因此速度不是很快（np.where是最慢的步骤）。

def setmass2(members, groups):
    idx = np.where(members['groupID'] == groups['groupID'][:,None])
    members['groupmass'][idx[1]] = groups['mass'][idx[0]]

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

2 个答案: