Skip to main content
 首页 » 编程设计

python之使用两个一维数组有效地索引二维 numpy 数组

2025年05月04日22myhome

我有一个大的 2d numpy 数组和两个 1d 数组,它们表示 2d 数组中的 x/y 索引。我想使用这些一维数组对二维数组执行操作。 我可以使用 for 循环来完成此操作,但在处理大型数组时速度非常慢。有没有更快的方法?我尝试将 1d 数组简单地用作索引,但这没有用。看这个例子:

import numpy as np 
 
# Two example 2d arrays 
cnt_a   =   np.zeros((4,4)) 
cnt_b   =   np.zeros((4,4)) 
 
# 1d arrays holding x and y indices 
xpos    =   [0,0,1,2,1,2,1,0,0,0,0,1,1,1,2,2,3] 
ypos    =   [3,2,1,1,3,0,1,0,0,1,2,1,2,3,3,2,0] 
 
# This method works, but is very slow for a large array 
for i in range(0,len(xpos)): 
    cnt_a[xpos[i],ypos[i]] = cnt_a[xpos[i],ypos[i]] + 1 
 
# This method is fast, but gives incorrect answer 
cnt_b[xpos,ypos] = cnt_b[xpos,ypos]+1 
 
 
# Print the results 
print 'Good:' 
print cnt_a 
print '' 
print 'Bad:' 
print cnt_b 

输出结果是:

Good: 
[[ 2.  1.  2.  1.] 
 [ 0.  3.  1.  2.] 
 [ 1.  1.  1.  1.] 
 [ 1.  0.  0.  0.]] 
 
Bad: 
[[ 1.  1.  1.  1.] 
 [ 0.  1.  1.  1.] 
 [ 1.  1.  1.  1.] 
 [ 1.  0.  0.  0.]] 

对于 cnt_b 数组,numpy 显然没有正确求和,但我不确定如何在不求助于用于计算 cnt_a 的(v.低效)for 循环的情况下解决这个问题。

请您参考如下方法:

另一种使用一维索引(由 @Shai 建议)的方法扩展到回答实际问题:

>>> out = np.zeros((4, 4)) 
>>> idx = np.ravel_multi_index((xpos, ypos), out.shape) # extract 1D indexes 
>>> x = np.bincount(idx, minlength=out.size) 
>>> out.flat += x 

np.bincount计算每个索引在 xpos, ypos 中出现的次数并将它们存储在 x 中.

或者,按照@Divakar 的建议:

>>> out.flat += np.bincount(idx, minlength=out.size)