Siraj S. Siraj S. - 2 months ago 9
Python Question

get desired element in numpy array of unequal row lengths, without using for loop

i have the below numpy array:

array([['apple','banana','orange'],
['car','bike','train','ship','plane','scooter'],
['red','purple']], dtype=object)


the individual rows in the array are of unequal length, I want to get the last element of each row. I can get this by running a for loop but I guess there could be more direct way to doing so. The closest (wrong solution) i have is arr[:][-1] which gives me element of last row and arr[np.arange(len(arr)),-1] which throws an 'IndexError' error.

My desired output is:

array([['orange','scooter','purple']], dtype=object)


I will appreciate any guidance. thank you.

Answer

using Pandas:

In [87]: a
Out[87]: array([['apple', 'banana', 'orange'], ['car', 'bike', 'train', 'ship', 'plane', 'scooter'], ['red', 'purple']], dtype=object)

In [88]: df = pd.DataFrame(a)

In [93]: df
Out[93]:
                                          0
0                   [apple, banana, orange]
1  [car, bike, train, ship, plane, scooter]
2                             [red, purple]

In [94]: df[0].str[-1]
Out[94]:
0     orange
1    scooter
2     purple
Name: 0, dtype: object

or as a NumPy array:

In [95]: df[0].str[-1].values
Out[95]: array(['orange', 'scooter', 'purple'], dtype=object)