Indexing one-dimensional numpy.array as matrix

I am trying to index a numpy.array with varying dimensions during runtime. To retrieve e.g. the first row of a n*m array a, you can simply do

a[0,:]

However, in case a happens to be a 1xn vector, this code above returns an index error:

IndexError: too many indices

As the code needs to be 开发者_运维知识库executed as efficiently as possible I don't want to introduce an if statement. Does anybody have a convenient solution that ideally doesn't involve changing any data structure types?

Just use a[0] instead of a[0,:]. It will return the first line for a matrix and the first entry for a vector. Is this what you are looking for?

If you want to get the whole vector in the one-dimensional case instead, you can use numpy.atleast_2d(a)[0]. It won't copy your vector -- it will just access it as a two-dimensional 1 x n-array.

From the 'array' or 'matrix'? Which should I use? section of the Numpy for Matlab Users wiki page:

For array, the vector shapes 1xN, Nx1, and N are all different things. Operations like A[:,1] return a rank-1 array of shape N, not a rank-2 of shape Nx1. Transpose on a rank-1 array does nothing.

Here's an example showing that they are not the same:

>>> import numpy as np
>>> a1 = np.array([1,2,3])
>>> a1
array([1, 2, 3])
>>> a2 = np.array([[1,2,3]])    // Notice the two sets of brackets
>>> a2
array([[1, 2, 3]])
>>> a3 = np.array([[1],[2],[3]])
>>> a3
array([[1],
       [2],
       [3]])

So, are you sure that all of your arrays are 2d arrays, or are some of them 1d arrays?

If you want to use your command of array[0,:], I would recommend actually using 1xN 2d arrays instead of 1d arrays. Here's an example:

>>> a2 = np.array([[1,2,3]])    // Notice the two sets of brackets
>>> a2
array([[1, 2, 3]])
>>> a2[0,:]
array([1, 2, 3])
>>> b2 = np.array([[1,2,3],[4,5,6]])
>>> b2
array([[1, 2, 3],
       [4, 5, 6]])
>>> b2[0,:]
array([1, 2, 3])