Search code examples
pythonarraysnumpyrowenumerate

Iterating over Numpy Array rows in python even with 1 row / 2 columns array


I am trying to iterate over the rows of a numpy array. The array may consist of two columns and multiple rows like [[a, b], [c, d], ...], or sometimes a single row like [a, b].

For the one-dimensional array, when I use enumerate to iterate over rows, python yields the individual elements a then b instead of the complete row [a, b] all at once.

How to I iterate the one-dimensional case in the same way as I would the 2D case?


Solution

  • Numpy iterates over the first dimension no matter what. Check the shape before you iterate.

    >>> x = np.array([1, 2])
    >>> x.ndim
    1
    
    >>> y = np.array([[1, 2], [3, 4], [5, 6]])
    >>> y.ndim
    2
    

    Probably the simplest method is to always wrap in a call to np.array:

    >>> x = np.array(x, ndmin=2, copy=False)
    >>> y = np.array(y, ndmin=2, copy=False)
    

    This will prepend a dimension of shape 1 to your array as necessary. It has the advantage that your inputs don't even have to be arrays, just something that can be converted to an array.

    Another option is to use the atleast_2d function:

    >>> x = np.atleast_2d(x)
    

    All that being said, you are likely sacrificing most of the benefits of using numpy in the first place by attempting a vanilla python loop. Try to vectorize your operation instead.