需要Numpy花哨的索引建議

Question

首先，對於不太完美的描述性標題感到抱歉。我有一個（500,500,500）ndarray「A」的浮點數，以及另一個（500,500,500）ndarray「Ind」整數。我想要形成第三個（500,500,500）ndarray「Anext」，其中「Anext」中索引（x，y，z）的條目是A [B [x，y，z]， y，z]。

我已經看過索引，但是我發現很難想出一個有用的解決方案。我感謝任何幫助。

Answer 1

使用numpy.indices()創建軸線1和軸線2的索引陣列，這裏是一個例子：

import numpy as np 

a = np.random.rand(4, 4, 4) 
ind = np.random.randint(0, 4, size=(4, 4, 4)) 

i0, i1, i2 = np.indices(a.shape) 

b = a[ind, i1, i2] 

查詢的結果：

for (i, j, k), v in np.ndenumerate(b): 
    assert v == a[ind[i, j, k], j, k] 

的i0形狀，i1，i2爲(4, 4, 4)，爲了節省一些內存，也可以使用numpy.ix_()：

i0, i1, i2 = np.ix_(*[np.arange(d) for d in a.shape]) 

其中i1.shape，i2.shape是(1, 4, 1), (1, 1, 4)