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

它正确地索引不同维度的数组。

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

1 个答案: