Question

考虑3d numpy数组：

ax1, ax2, ax3 = 3, 3, 2
arr = np.asarray(range(ax1*ax2*ax3)).reshape([ax1, ax2, ax3])
arr: [[[0, 1], [2, 3], [4, 5]],
      [[6, 7], [8, 9], [10, 11]],
      [[12, 13], [14, 15], [16, 17]]]

和索引向量idx = [0, 1, 2]。

我想通过以下语句将数组arr与idx切片：

res = [arr[i, :idx[i]+1] for i in range(ax1)]
res: [[[0, 1]],
      [[6, 7], [8, 9],
      [[12, 13], [14, 15], [16, 17]]]

但这种切片看起来很复杂。

numpy是否支持这样的操作而不使用循环？我正在寻找类似arr[range(ax1), :idx+1]的内容。

Answer 1

您的问题是结果值不矩形：您不能将其正确表示为数组。

如果您只使用不同格式的值，那么您可以通过boolean mask获得所需内容：

>>> mask = np.tri(3, 3, dtype=bool)
>>> arr[mask]
array([[ 0,  1],
       [ 6,  7],
       [ 8,  9],
       [12, 13],
       [14, 15],
       [16, 17]])

原则是你为[0; 2] ^ 2中的每对索引传递，是否应该使用该对：

>>> np.tri(3, 3, dtype=bool)
array([[ True, False, False],
       [ True,  True, False],
       [ True,  True,  True]], dtype=bool)

这导致了非常简洁：

>>> arr[np.tri(3, 3, dtype=bool)]
array([[ 0,  1],
       [ 6,  7],
       [ 8,  9],
       [12, 13],
       [14, 15],
       [16, 17]])

Answer 2

这是一种假设arr和idx作为NumPy数组的矢量化方法 -

np.split(arr[np.arange(arr.shape[0]) <= idx[:,None]],(idx+1).cumsum())[:-1]

运行样本以验证结果 -

In [5]: arr
Out[5]: 
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5]],

       [[ 6,  7],
        [ 8,  9],
        [10, 11]],

       [[12, 13],
        [14, 15],
        [16, 17]]])

In [6]: idx
Out[6]: array([2, 0, 1])

In [7]: np.split(arr[np.arange(arr.shape[0]) <= idx[:,None]],(idx+1).cumsum())[:-1]
Out[7]: 
[array([[0, 1],
        [2, 3],
        [4, 5]]), array([[6, 7]]), array([[12, 13],
        [14, 15]])]

In [8]: [arr[i, :idx[i]+1] for i in range(ax1)] # Loopy approach
Out[8]: 
[array([[0, 1],
        [2, 3],
        [4, 5]]), array([[6, 7]]), array([[12, 13],
        [14, 15]])]

如何用索引向量切片3d ndarray

2 个答案: