Numpy: get 1D array from a 2D array with 1 row

Question

I have a complex function that returns a 2D array, to make it simple, let's consider the following:

import numpy as np

def get_array():
    a = np.array([[0, 9, 5]])
    return a

Is there a numpy command that allows me to retrieve the single row automatically instead of doing the following?

def get_array():
    a = np.array([[0, 9, 5]])
    if a.shape[0] == 1:
        return a[0]
    else:
        return a 
    return a

Thanks a lot!

a.ravel() or a.flatten()...

Quang Hoang
– Quang Hoang

2021-11-30 16:56:12 +00:00
Commented Nov 30, 2021 at 16:56 — Quang Hoang
– Quang Hoang, Commented Nov 30, 2021 at 16:56

orlp · Accepted Answer · 2021-11-30 16:56:38Z

2

np.squeeze(a) will remove any unit sized axes from a.

>>> np.squeeze([1, 2, 3])
array([1, 2, 3])

>>> np.squeeze([[1, 2, 3]])
array([1, 2, 3])

>>> np.squeeze([[1, 2, 3], [4, 5, 6]])
array([[1, 2, 3],
       [4, 5, 6]])

It might not be exactly what you want though, as also:

>>> np.squeeze([[1], [2]])
array([1, 2])

answered Nov 30, 2021 at 16:56

orlp

119k39 gold badges226 silver badges324 bronze badges

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

2 Comments

mauro Over a year ago

Thank you @orlp, this is exactly what I need as in my case the 2D array is (Nx3)

orlp Over a year ago

@mauro If my answer solved your problem please consider marking it as accepted.

Learning is a mess · Accepted Answer · 2021-11-30 17:05:28Z

0

On top of @orlp answer, if you are conscious of performance, ravel seems to be faster than squeeze:

answered Nov 30, 2021 at 17:05

Learning is a mess

8,3178 gold badges42 silver badges82 bronze badges

Collectives™ on Stack Overflow

Numpy: get 1D array from a 2D array with 1 row

2 Answers 2

2 Comments

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

2 Comments

Comments

Your Answer

Sign up or log in

Post as a guest

Related