我有一個 3D numpy 陣列 shape (i, j, k)。我有一個長度陣列，i其中包含k. 我想索引陣列以獲得一個形狀(i, j)。
這是我試圖實作的一個例子：
import numpy as np
arr = np.arange(2 * 3 * 4).reshape(2, 3, 4)
# 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]]])
indices = np.array([1, 3])
# I want to mask `arr` using `indices`
# Desired output is equivalent to
# np.stack((arr[0, :, 1], arr[1, :, 3]))
# array([[ 1, 5, 9],
# [15, 19, 23]])
我嘗試重塑indices陣列以使其能夠進行廣播，arr但這會引發IndexError.
arr[indices[np.newaxis, np.newaxis, :]]
# IndexError: index 3 is out of bounds for axis 0 with size 2
我還嘗試創建一個 3D 蒙版并將其應用于arr. 這對我來說似乎更接近正確答案，但我仍然以IndexError.
mask = np.stack((np.arange(arr.shape[0]), indices), axis=1)
arr[mask.reshape(2, 1, 2)]
# IndexError: index 3 is out of bounds for axis 0 with size 2
uj5u.com熱心網友回復：
根據我在您的示例中的理解，您可以簡單地將indices作為您的第二維切片傳遞，并且range長度與您的第零維切片的索引相對應，如下所示：
import numpy as np
arr = np.arange(2 * 3 * 4).reshape(2, 3, 4)
indices = np.array([1, 3])
print(arr[range(len(indices)), :, indices])
# array([[ 1, 5, 9],
# [15, 19, 23]])
uj5u.com熱心網友回復：
這有效：
sub = arr[[0,1], :, [1,3]]
輸出：
>>> sub
array([[ 1, 5, 9],
[15, 19, 23]])
@Psidom 的一個更動態的版本：
>>> sub = arr[np.arange(len(arr)), :, [1,3]]
array([[ 1, 5, 9],
[15, 19, 23]])
轉載請註明出處，本文鏈接：https://www.uj5u.com/yidong/389392.html