user2489252
user2489252

Reputation:

Best way to iterate through a numpy array returning the columns as 2d arrays

EDIT: Thank you all for the good solutions, I think if I'd had to pick one, it would be A[:,[0]]

I collected 7 approaches now and put them into an IPython notebook. The timeit benchmarks are not suprising: they are all roughly the same in terms of speed.

Thanks a lot for your suggestion!


I a looking for a good way to iterate through the columns of a matrix and return them as 1xd column vectors. I have some ideas, but I don't think that those are good solutions. I think I am missing something here. Which way would you recommend? E.g., let's say I have the following matrix and want to return the first column as a column vector:

A = np.array([ [1,2,3], [4,5,6], [7,8,9] ])
>>> A
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])

By default, numpy returns it like this:

>>> A[:,0]
array([1, 4, 7])
>>> A[:,0].shape
(3,)

And what I want is this:

array([[1],
       [4],
       [7]])

with .shape = (3,1)


Transpose doesn't work to return it as a column vector.

>>> A[:,0].T
array([1, 4, 7])

>>> A[:,0]
  array([1, 4, 7])

I would have to create a new axis every time

>>> A[:,0][:,np.newaxis].shape
(3, 1)
>>> A[:,0][:,np.newaxis]
array([[1],
       [4],
       [7]])

Or after doing some experimenting, I came up with other workarounds like this:

>>> A[:,0:1]
array([[1],
       [4],
       [7]])

>>> A[:,0].reshape(A.shape[1],1)
array([[1],
       [4],
       [7]])

Upvotes: 0

Views: 239

Answers (3)

Kiwi
Kiwi

Reputation: 2816

My favorite solution is the slicing. You have different solutions :

A[:,0:1] # not so clear
A[:,:1]  # black magic
A[:,[0]] # clearest syntax imho

Concerning the reshape solution, you can enhance the syntax like this :

A[:,0].reshape(A.shape[1],1)
A[:,0].reshape(-1,1)

You can also merge the following :

A[:,0][:,np.newaxis] # ->
A[:,0,np.newaxis] # or
A[:,np.newaxis,0]

Upvotes: 1

dawg
dawg

Reputation: 103844

You can use column_stack:

>>> A
array([[1, 2, 3],
       [4, 5, 6],
       [7, 8, 9]])
>>> np.column_stack((A[:,0],))
array([[1],
       [4],
       [7]])
>>> #                ^^^^^^^ a tuple

Just make sure that you are feeding it a 1 element tuple for a single column or you are getting something different:

>>> np.column_stack(A[:,0])
array([[1, 4, 7]])

Upvotes: 0

Ashwini Chaudhary
Ashwini Chaudhary

Reputation: 250961

One way would be to use numpy.row_stack or numpy.vstack:

In [91]: np.row_stack(A[:,0])
Out[91]: 
array([[1],
       [4],
       [7]])
In [92]: np.vstack(A[:,0])
Out[92]: 
array([[1],
       [4],
       [7]])

Upvotes: 0

Related Questions