在numpy数组中沿轴获取最后一个元素

Question

我需要一个函数来获取 numpy 数组中沿轴的最后一个元素。

例如，如果我有一个数组，

a = np.array([1, 2, 3])

函数应该像这样工作

get_last_elements(a, axis=0)
>>> [3]
get_last_elements(a, axis=1)
>>> [1, 2, 3]

这个函数也需要对多维数组起作用：

b = np.array([[1, 2],
              [3, 4]])

get_last_elements(b, axis=0)
>>> [[2],
     [4]]
get_last_elements(b, axis=1)
>>> [3, 4]

有人有实现它的好主意吗？

Answer 1

您可以使用np.take 来获取：

def get_last_elements(a, axis=0):
  shape = list(a.shape)
  shape[axis] = 1
  return np.take(a,-1,axis=axis).reshape(tuple(shape))

输出：

print(get_last_elements(b, axis=0))
[[3 4]]

print(get_last_elements(b, axis=1))
[[2]
 [4]]