python - 根据用户输入返回多个矩阵到函数

Question

我有一个接受用户参数 k 并生成以下形式的矩阵的函数:

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

离开函数我想返回 k 2x5 矩阵。本质上，

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