N dimensional arrays - Python/Numpy

Question

just wondering if there is any clever way to do the following.

I have an N dimensional array representing a 3x3 grid

grid = [[1,2,3],
        [4,5,6],
        [7,8,9]]

In order to get the first row I do the following:

grid[0][0:3]
>> [1,2,3]

In order to get the first column I would like to do something like this (even though it is not possible):

grid[0:3][0]
>> [1,4,7]

Does NumPy support anything similar to this by chance?

---

Any ideas?

Answer 1

Yes, there is something like that in Numpy:

import numpy as np

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

grid[0,:]
# array([1, 2, 3])

grid[:,0]
# array([1, 4, 7])

Answer 2

You can use zip to transpose a matrix represented as a list of lists:

>>> zip(*grid)[0]
(1, 4, 7)

Anything more than just that, and I'd use Numpy.

Answer 3

To get the columns in Python you could use:

[row[0] for row in grid]
>>> [1,4,7]

You could rewrite your code for getting the row as

grid[0][:]

because [:] just copies the whole array, no need to add the indices.

However, depending on what you want to achieve, I'd say it's better to just write a small matrix class to hide this implementation stuff.