不对称的切片蟒蛇

Question

考虑以下矩阵:

X = np.arange(9).reshape(3,3)
     array([[0, 1, 2],
            [3, 4, 5],
            [6, 7, 8]]) 

假设我想要对以下数组进行子集化

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

例如,可以对行和列进行一些索引

col=[0,1,2] 
row = [[0,1],[1,2],[0,1]]

然后,如果我将结果存储在变量数组中,我可以使用以下代码执行此操作:

array=np.zeros([2,3],dtype='int64')
for i in range(3):
    array[:,i]=X[row[i],col[i]]

有没有办法播放这种手术？我必须将此作为大型文件〜5 Gb的数据清理阶段,我想使用dask来并行化它.但是,如果我第一次避免使用for循环,我会感觉很棒.

Answer 1

对于带有 的数组NumPy's advanced-indexing，它将是 -

X[row, np.asarray(col)[:,None]].T

样本运行 -

In [9]: X
Out[9]: 
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

In [10]: col=[0,1,2] 
    ...: row = [[0,1],[1,2],[0,1]]

In [11]: X[row, np.asarray(col)[:,None]].T
Out[11]: 
array([[0, 4, 2],
       [3, 7, 5]])