如何在python中指定未知数量的前面的索引？

Question

我将(n_features, 3)或(n_samples, n_features, 3)保存在数组中。你知道吗

有没有可能在不知道前面维度的数目的情况下，从结尾含蓄地处理第二维度？你知道吗

现在我写了

def get_feature(data, i):
   n = data.ndim
   if n == 2:
      return data[i,:]
   elif n == 3:
      return data[:,i,:]
   elif n == 4:
      return data[:,:,i,:]
   else:
      raise NotImplementedError

如何缩短？你知道吗

Answer 1

可以使用Ellipsis跳过前几个维度：

data[...,i,:]

---

示例：

a = np.arange(24).reshape(2,2,2,3)
b = np.arange(24).reshape(2,4,3)
c = np.arange(24).reshape(8,3)

(a[...,1,:] == a[:,:,1,:]).all()
# True

(b[...,1,:] == b[:,1,:]).all()
# True

(c[...,1,:] == c[1,:]).all()
# True

它可以正确地索引不同维度的数组。你知道吗