Daniel Keast

Python array columns

programming, python

I wanted to access the columns in a 2D array in Python today. I found that it’s surprisingly simple to do. This is what I ended up with:

>>> li = [[1, 2], [3, 4], [5, 6]]
>>> x, y = zip(*li)
>>> x
(1, 3, 5)
>>> y
(2, 4, 6)
>>>

The asterisk in the zip function call is the argument unpacking feature of the python call syntax. It explodes an iterator into positional arguments for the function, which in this case means it is removing the outer array from li and passing the inner three to the zip function. It is the equivalent of this:

>>> li = [[1, 2], [3, 4], [5, 6]]
>>> one = li[0]
>>> two = li[1]
>>> three = li[2]
>>> zip(one, two, three)

This will fail if the iterable does not contain exactly the correct number of elements, which feels like the correct thing to be doing:

>>> def f(a, b):
...    return a + b
...
>>> f(*[1, 2])
3
>>> f(*[1])
Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
TypeError: f() missing 1 required positional argument: 'b'
>>> f(*[1, 2, 3])
Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
TypeError: f() takes 2 positional arguments but 3 were given
>>>                                               

The other interesting part is the x, y = syntax. The zip function returns an iterator of tuples. This means you can use the iterator unpacking feature to destructure the response into x and y axis tuples containing the columns of the li array.