沿着最後維度的索引numpy nd數組

有沒有簡單的方法來使用索引數組索引沿最後維度的numpy多維數組？例如，形狀爲(10, 10, 20)的數組a。假設我有一個形狀爲(10, 10)的索引b的數組，因此結果將爲c[i, j] = a[i, j, b[i, j]]。沿着最後維度的索引numpy nd數組

我試過下面的例子：

a = np.ones((10, 10, 20)) 
b = np.tile(np.arange(10) + 10, (10, 1)) 
c = a[b]

但是，這並不工作，因爲它然後嘗試指數像a[b[i, j], b[i, j]]，這是不一樣的a[i, j, b[i, j]]。等等。有沒有一個簡單的方法來做到這一點，而不訴諸於循環？

來源

2014-12-03 tiago

只是爲了確保我的理解正確，你想'C [I，J] = A [1， j，b [i，j]]' – mgilson 2014-12-03 17:00:18

@mgilson是的。應該做得更清楚。 – tiago 2014-12-03 18:19:41

這通常是'c = a [np.arange（b.shape [0]），np.arange（b.shape [1]），b]'完成的，但我希望有更好的方法。 – 2014-12-03 20:31:48

有幾種方法可以做到這一點。讓我們先來生成一些測試數據：要解決的問題是要建立兩個指標向量，其中一個是行向量和0..9使用meshgrid其他的列向量

In [1]: a = np.random.rand(10, 10, 20) 

In [2]: b = np.random.randint(20, size=(10,10)) # random integers in range 0..19

方式一：

In [3]: i1, i0 = np.meshgrid(range(10), range(10), sparse=True) 

In [4]: c = a[i0, i1, b]

這是可行的，因爲i0,i1和b都將被廣播到10x10矩陣。正確性快速測試：

In [5]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10)) 
Out[5]: True

另一種方法是使用choose和rollaxis：

# choose needs a sequence of length 20, so move last axis to front 
In [22]: aa = np.rollaxis(a, -1) 

In [23]: c = np.choose(b, aa) 

In [24]: all(c[i, j] == a[i, j, b[i, j]] for i in range(10) for j in range(10)) 
Out[24]: True

來源

2014-12-03 21:01:14

很好地完成。我喜歡。 – mgilson 2014-12-03 22:34:10

謝謝。我更喜歡你的第一個選擇。 'choose'選項看起來不錯，但它不夠普遍，因爲'choose'不適用於大於32的索引（當前的numpy錯誤，請參閱[關於增加'NPY_MAXARGS'的討論]（https://github.com/numpy的/ numpy的/問題/ 4398））。根據我的測試，它也是'meshgrid'選項的兩倍。 – tiago 2014-12-04 09:45:22

在這種情況下，您應該查看'mgrid'和'ogrid'等函數，這些函數是從'meshgrid'文檔中引用的。但是你使用32維以上的矩陣嗎？我爲你感到難過，當試圖理解超過2維的陣列時，我的頭已經爆炸了:)。 – 2014-12-04 09:53:58

沿着最後維度的索引numpy nd數組

回答

相關問題