numpy切片返回最后兩個維度

Question

基本上，我正在尋找一種函數或語法，該函數或語法將允許我獲取具有任意數量維的維numpy數組最后兩個維的第一個“切片”。

我可以這樣做，但是很難忍受，如果有人發送6d數組怎么辦？ 必須有一個像橢圓這樣的numpy函數，它可以擴展為0,0,0，...而不是：，：，：，...

data_2d = np.ones(5**2).reshape(5,5)
data_3d = np.ones(5**3).reshape(5,5,5)
data_4d = np.ones(5**4).reshape(5,5,5,5)

def get_last2d(data):
    if data.ndim == 2:
        return data[:]
    if data.ndim == 3:
        return data[0, :]
    if data.ndim == 4:
        return data[0, 0, :]

np.array_equal(get_last2d(data_3d), get_last2d(data_4d))

謝謝，科林

Answer 1

這個怎么樣，

def get_last2d(data):
    if data.ndim <= 2:
        return data
    slc = [0] * (data.ndim - 2)
    slc += [slice(None), slice(None)]
    return data[slc]

Answer 2

def get_last_2d(x):
    m,n = x.shape[-2:]
    return x.flat[:m*n].reshape(m,n)

之所以可行，是因為對數組進行展平將按快速變化的索引順序返回條目，並且對於C風格的索引，最后的索引變化最快。 因此，平整數組的前m * n個條目就是您想要的。