Wie ein 3D ndarray mit einem Indexvektor

Question

Betrachtet man einen 3D numpy Array schneiden:

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]]] 

und einen Indexvektor idx = [0, 1, 2].

Ich möchte das Array arr mit idx durch die folgende Aussage schneiden:

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

Aber diese Art des Schneidens kompliziert aussieht.

Unterstützt numpy einen solchen Vorgang ohne Schleife? Ich suche nach etwas wie arr[range(ax1), :idx+1].

Answer 1

Ihr Problem ist, dass die resultierenden Werte sind nicht rechteckig: Sie kann das nicht richtig als Array darstellen.

Wenn Sie mit mit den Werte fein sind nur in einem anderen Format, können Sie bekommen, was Sie brauchen einen durch boolean mask:

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

Das Prinzip ist, dass Sie passieren, für jedes Paar von Indizes in [0; 2]^2, ob Sie dieses Paar oder nicht nehmen sollen:

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

, die zu dem wunderbar prägnanten führt:

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

Answer 2

Hier ist ein vektorisiert Ansatz arr und idx als NumPy Arrays vorausgesetzt -

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

Beispiellaufergebnisse zu verifizieren -

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]])]