返回基於用戶輸入的多個矩陣以功能

Question

我有接受用戶參數k，併產生以下形式的矩陣的函數：

a = np.array((k,2,5)) 

離開我想返回ķ 2×5矩陣的功能。從本質上講，

if k == 2: 
    return a[0,:,:], a[1,:,:] 
if k == 3: 
    return a[0,:,:], a[1,:,:], a[2,:,:] 

有沒有一種辦法，而不必硬編碼每個k的可能值的動態做到這一點？

謝謝。

Answer 1

肯定的：

return tuple(a[i,:,:] for i in range(k)) 

Python的偉大！ 8-））

Answer 2

下面是一個使用np.array_split的方法 -

np.array_split(a[:k].reshape(-1,a.shape[-1]),k)shorter 

即使np.vsplit短 -

np.vsplit(a[:k].reshape(-1,a.shape[-1]), k) 

---

請注意，如果您正在分配從相關函數輸出k輸出的數量，你不需要所有這些混亂，並可以簡單地返回片。即：return a[:k]。

作爲一個例子：

In [595]: a = np.random.rand(4,2,5) 

In [596]: k = 3 

In [597]: a0,a1,a2 = a[:k] 

In [598]: np.allclose(a0, a[0,:,:]) # Verify for all three outputs 
Out[598]: True 

In [599]: np.allclose(a1, a[1,:,:]) 
Out[599]: True 

In [600]: np.allclose(a2, a[2,:,:]) 
Out[600]: True