python – 访问数组的多个元素

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.4,1.3,1.5,2.4,2.6,3.5,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,4],4,5,6,7]])

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

col_start = np.array([2,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])[:,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])[:,idx]
10000 loops,best of 3: 108 μs per loop

In [537]: timeit A.flat[idx+np.arange(A.shape[0])[:,best of 3: 59.4 μs per loop

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

什么超出idx？

col_start=np.array([2,8])
idx=col_start[:,None]+np.arange(3)
A[np.arange(A.shape[0])[:,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])[:,idx]

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

startend = [[2,[6,9],[8,10]]
index_list = [np.arange(v[0],array([14,15,16]),array([26,27,28]),array([38,39])]

A.flat[np.r_[tuple(index_list)]]
# array([ 0.2,1.6,2.7,2.8,3.8,3.9])