通过另一个数组有效地索引多维numpy数组

时间:2016-07-22 22:01:58

标签: python arrays numpy

我有一个数组x,我想要访问哪些特定值,其索引由另一个数组给出。

例如,x

array([[ 0,  1,  2,  3,  4],
[ 5,  6,  7,  8,  9],
[10, 11, 12, 13, 14],
[15, 16, 17, 18, 19],
[20, 21, 22, 23, 24]])

并且索引是Nx2

的数组
idxs = np.array([[1,2], [4,3], [3,3]])

我想要一个返回x [1,2],x [4,3],x [3,3]或[7,23,18]数组的函数。下面的代码可以解决这个问题,但是我希望通过避免for循环来加速大型数组。

import numpy as np

def arrayvalsofinterest(x, idx):
    output = np.zeros(idx.shape[0])
    for i in range(len(output)):
        output[i] = x[tuple(idx[i,:])]
    return output

if __name__ == "__main__":
    xx = np.arange(25).reshape(5,5)
    idxs = np.array([[1,2],[4,3], [3,3]])
    print arrayvalsofinterest(xx, idxs)

1 个答案:

答案 0 :(得分:3)

您可以传递一个可迭代的axis0坐标和一个可迭代的axis1坐标。请参阅the Numpy docs here

i0, i1 = zip(*idxs)
x[i0, i1]

正如@Divakar在评论中指出的那样,这比使用数组视图的内存效率更低,即

x[idxs[:, 0], idxs[:, 1]]