python – 访问数组的多个元素

2023年1月11日 178次阅读

有没有办法在一个操作中为这些元素的已知行和列获取数组元素？在每一行中,我想访问从col_start到col_end的元素(每行有不同的起始和结束索引).每行的元素数量相同,元素是连续的.

例：

[ . . . . | | | . . . . . ]
[ | | | . . . . . . . . . ]
[ . . | | | . . . . . . . ]
[ . . . . . . . . | | | . ]

一种解决方案是获取元素的索引(行 – 列对),而不是使用my_array [row_list,col_list].

没有使用for循环有没有其他(更简单)的方法？

最佳答案

A = np.arange(40).reshape(4,10)*.1
startend = [[2,5],[3,6],[4,7],[5,8]]
index_list = [np.arange(v[0],v[1]) + i*A.shape[1] 
                 for i,v in enumerate(startend)]
# [array([2, 3, 4]), array([13, 14, 15]), array([24, 25, 26]), array([35, 36, 37])]
A.flat[index_list]

生产

array([[ 0.2,  0.3,  0.4],
       [ 1.3,  1.4,  1.5],
       [ 2.4,  2.5,  2.6],
       [ 3.5,  3.6,  3.7]])

这仍然有一个迭代,但它是一个相当基本的列表.
我正在索引扁平的1d版本A. np.take(A,index_list)也可以.

如果行间隔大小不同,我可以使用np.r_来连接它们.这不是绝对必要的,但是在从多个区间和值构建索引时这是一种方便.

A.flat[np.r_[tuple(index_list)]]
# array([ 0.2,  0.3,  0.4,  1.3,  1.4,  1.5,  2.4,  2.5,  2.6,  3.5,  3.6, 3.7])

ajcr使用的idx可以在不选择的情况下使用：

idx = [np.arange(v[0], v[1]) for i,v in enumerate(startend)]
A[np.arange(A.shape[0])[:,None], idx]

idx就像我的index_list,除了它不添加行长度.

np.array(idx)

array([[2, 3, 4],
       [3, 4, 5],
       [4, 5, 6],
       [5, 6, 7]])

由于每个arange具有相同的长度,因此可以在不迭代的情况下生成idx：

col_start = np.array([2,3,4,5])
idx = col_start[:,None] + np.arange(3)

第一个索引是一个列数组,广播以匹配此idx.

np.arange(A.shape[0])[:,None] 
array([[0],
       [1],
       [2],
       [3]])

有了这个A和idx,我得到以下时间：

In [515]: timeit np.choose(idx,A.T[:,:,None])
10000 loops, best of 3: 30.8 µs per loop

In [516]: timeit A[np.arange(A.shape[0])[:,None],idx]
100000 loops, best of 3: 10.8 µs per loop

In [517]: timeit A.flat[idx+np.arange(A.shape[0])[:,None]*A.shape[1]]
10000 loops, best of 3: 24.9 µs per loop

平面索引更快,但计算发烧友指数需要一些时间.

对于大型阵列,平面索引的速度占主导地位.

A=np.arange(4000).reshape(40,100)*.1
col_start=np.arange(20,60)
idx=col_start[:,None]+np.arange(30)

In [536]: timeit A[np.arange(A.shape[0])[:,None],idx]
10000 loops, best of 3: 108 µs per loop

In [537]: timeit A.flat[idx+np.arange(A.shape[0])[:,None]*A.shape[1]]
10000 loops, best of 3: 59.4 µs per loop

np.choose方法遇到硬编码限制：需要介于2和(32)个数组对象之间(包括).

什么超出idx？

col_start=np.array([2,4,6,8])
idx=col_start[:,None]+np.arange(3)
A[np.arange(A.shape[0])[:,None], idx]

产生错误,因为最后一个idx值是10,太大了.

你可以剪辑idx

idx=idx.clip(0,A.shape[1]-1)

在最后一行中生成重复值

[ 3.8,  3.9,  3.9]

您也可以在编制索引之前填充A.有关更多选项,请参阅np.pad.

np.pad(A,((0,0),(0,2)),'edge')[np.arange(A.shape[0])[:,None], idx]

另一种选择是删除越界值.然后,idx将成为列表(或列表数组)的参差不齐的列表.平面方法可以处理这个问题,但结果不是矩阵.

startend = [[2,5],[4,7],[6,9],[8,10]]
index_list = [np.arange(v[0],v[1]) + i*A.shape[1] 
                 for i,v in enumerate(startend)]
# [array([2, 3, 4]), array([14, 15, 16]), array([26, 27, 28]), array([38, 39])]

A.flat[np.r_[tuple(index_list)]]
# array([ 0.2,  0.3,  0.4,  1.4,  1.5,  1.6,  2.6,  2.7,  2.8,  3.8,  3.9])
点赞