我需要一個函數來獲取 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]有沒有人有實現它的好主意？
1 回答
Helenr
TA貢獻1780條經驗 獲得超4個贊
您可以使用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]]
添加回答
舉報
0/150
提交
取消