Getting last elements along an axis in numpy array

Question

I need a function to get last elements along an axis in numpy array.

For example, If I have a array,

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

The function should work like

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

This function needs to work for multidimensional array too:

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

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

Does anyone have a good idea to achieve it?

Are you sure you want get_last_elements(a, axis=1) to work for the first example? Since a is of shape (3,), not (1, 3). — Mario Ishac
– Mario Ishac, Commented Aug 1, 2020 at 11:21
Yes. I want that. But, maybe I can convert 1d array to 2d with X = np.reshape(X,(1, X.size)) or something easily. — Atsushi Sakai
– Atsushi Sakai, Commented Aug 1, 2020 at 11:25
np.take may be useful. np.atleast_2d may help when dealing with a 1d input. — hpaulj
– hpaulj, Commented Aug 1, 2020 at 15:38

Ehsan · Accepted Answer · 2020-08-02 14:51:22Z

1

You can use np.take to get that:

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

output:

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

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

answered Aug 1, 2020 at 22:44

Ehsan

12.5k2 gold badges24 silver badges36 bronze badges

Sign up to request clarification or add additional context in comments.

1 Answer 1