在numpy ndarray中重新排序和重塑元素

Question

我有一個由外部模塊返回的numpy ndarray。數組的形狀是（3,3,128）。基本上是一堆128塊瓷磚，每塊瓷磚都是3x3。

如何重新排列形狀使其變成（128,3,3）。這樣就可以更容易地通過圖塊編號進行索引。最後一步是平坦化到（128,9），這樣每個128塊可以很容易地作爲一個9值向量來訪問。

Answer 1

您可以使用指定的，例如新的數組順序調換，

a = np.arange(0,3*3*128).reshape(3,3,128) 
a_reorder = a.transpose([2,0,1]) 

，你可以檢查是通過比較所有的瓷磚正確，

np.all([np.all(a[:,:,i]==a_reorder[i,:,:]) for i in range(128)]) 

，並壓平

a_flat = a_reorder.reshape(128,9)