What's the best way to access columns of an array in Python?
Question:
In Matlab, one can access a column of an array with ::
>> array=[1 2 3; 4 5 6]
array =
1 2 3
4 5 6
>> array(:,2)
ans =
2
5
How to do this in Python?
>>> array=[[1,2,3],[4,5,6]]
>>> array[:,2]
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: list indices must be integers, not tuple
>>> array[:][2]
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
Addendum
I’d like an example applied to an array of dimensions greater than three:
>> B = cat(3, eye(3), ones(3), magic(3))
B(:,:,1) =
1 0 0
0 1 0
0 0 1
B(:,:,2) =
1 1 1
1 1 1
1 1 1
B(:,:,3) =
8 1 6
3 5 7
4 9 2
>> B(:,:,1)
ans =
1 0 0
0 1 0
0 0 1
>> B(:,2,:)
ans(:,:,1) =
0
1
0
ans(:,:,2) =
1
1
1
ans(:,:,3) =
1
5
9
Answers:
def get_column(array, col):
result = []
for row in array:
result.appen(row[col])
return result
Use like this (remember that indexes start from 0):
>>> a = [[1,2,3], [2,3,4]]
>>> get_column(a, 1)
[2, 3]
Indexing / slicing with Python using the colon results in things a bit differently than matlab. If you have your array, [:] will copy it. If you want all values at a specific index of nested arrays, you probably want something like this:
array = [[1,2,3],[4,5,6]]
col1 = [inner[0] for inner in array] # note column1 is index 0 in Python.
Use a list comprehension to build a list of values from that column:
def nthcolumn(n, matrix):
return [row[n] for row in matrix]
Optionally use itemgetter if you need a (probably slight) performance boost:
from operator import itemgetter
def nthcolumn(n, matrix):
nthvalue = itemgetter(n)
return [nthvalue(row) for row in matrix]
If using nested lists, you can use a list comprehension:
array = [ [1, 2, 3], [4, 5, 6] ]
col2 = [ row[1] for row in array ]
Keep in mind that since Python doesn’t natively know about matrices, col2 is a list, and as such both “rows” and “columns” are the same type, namely lists. Use the numpy package for better support for matrix math.
You can group data in a two-dimensional list by column using the built-in zip() function:
>>> array=[[1,2,3],[4,5,6]]
>>> zip(*array)
[(1, 4), (2, 5), (3, 6)]
>>> zip(*array)[1]
(2, 5)
Note that the index starts at 0, so to get the second column as in your example you use zip(*array)[1] instead of zip(*array)[2]. zip() returns tuples instead of lists, depending on how you are using it this may not be a problem, but if you need lists you can always do map(list, zip(*array)) or list(zip(*array)[1]) to do the conversion.
If you use Matlab, you probably will want to install NumPy:
Using NumPy, you can do this:
In [172]: import numpy as np
In [173]: arr = np.matrix('1 2 3; 4 5 6')
In [174]: arr
Out[174]:
matrix([[1, 2, 3],
[4, 5, 6]])
In [175]: arr[:,2]
Out[175]:
matrix([[3],
[6]])
Since Python uses 0-based indexing (while Matlab uses 1-based indexing), to get the same slice you posted you would do:
In [176]: arr[:,1]
Out[176]:
matrix([[2],
[5]])
It is easy to build numpy arrays of higher dimension as well. You could use np.dstack for instance:
In [199]: B = np.dstack( (np.eye(3), np.ones((3,3)), np.arange(9).reshape(3,3)) )
In [200]: B.shape
Out[200]: (3, 3, 3)
In [201]: B[:,:,0]
Out[201]:
array([[ 1., 0., 0.],
[ 0., 1., 0.],
[ 0., 0., 1.]])
In [202]: B[:,:,1]
Out[202]:
array([[ 1., 1., 1.],
[ 1., 1., 1.],
[ 1., 1., 1.]])
In [203]: B[:,:,2]
Out[203]:
array([[ 0., 1., 2.],
[ 3., 4., 5.],
[ 6., 7., 8.]])
And here is the array formed from the second column from each of the 3 arrays above:
In [204]: B[:,1,:]
Out[204]:
array([[ 0., 1., 1.],
[ 1., 1., 4.],
[ 0., 1., 7.]])
Numpy doesn’t have a function to create magic squares, however. sniff
In Matlab, one can access a column of an array with ::
>> array=[1 2 3; 4 5 6]
array =
1 2 3
4 5 6
>> array(:,2)
ans =
2
5
How to do this in Python?
>>> array=[[1,2,3],[4,5,6]]
>>> array[:,2]
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
TypeError: list indices must be integers, not tuple
>>> array[:][2]
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
Addendum
I’d like an example applied to an array of dimensions greater than three:
>> B = cat(3, eye(3), ones(3), magic(3))
B(:,:,1) =
1 0 0
0 1 0
0 0 1
B(:,:,2) =
1 1 1
1 1 1
1 1 1
B(:,:,3) =
8 1 6
3 5 7
4 9 2
>> B(:,:,1)
ans =
1 0 0
0 1 0
0 0 1
>> B(:,2,:)
ans(:,:,1) =
0
1
0
ans(:,:,2) =
1
1
1
ans(:,:,3) =
1
5
9
def get_column(array, col):
result = []
for row in array:
result.appen(row[col])
return result
Use like this (remember that indexes start from 0):
>>> a = [[1,2,3], [2,3,4]]
>>> get_column(a, 1)
[2, 3]
Indexing / slicing with Python using the colon results in things a bit differently than matlab. If you have your array, [:] will copy it. If you want all values at a specific index of nested arrays, you probably want something like this:
array = [[1,2,3],[4,5,6]]
col1 = [inner[0] for inner in array] # note column1 is index 0 in Python.
Use a list comprehension to build a list of values from that column:
def nthcolumn(n, matrix):
return [row[n] for row in matrix]
Optionally use itemgetter if you need a (probably slight) performance boost:
from operator import itemgetter
def nthcolumn(n, matrix):
nthvalue = itemgetter(n)
return [nthvalue(row) for row in matrix]
If using nested lists, you can use a list comprehension:
array = [ [1, 2, 3], [4, 5, 6] ]
col2 = [ row[1] for row in array ]
Keep in mind that since Python doesn’t natively know about matrices, col2 is a list, and as such both “rows” and “columns” are the same type, namely lists. Use the numpy package for better support for matrix math.
You can group data in a two-dimensional list by column using the built-in zip() function:
>>> array=[[1,2,3],[4,5,6]]
>>> zip(*array)
[(1, 4), (2, 5), (3, 6)]
>>> zip(*array)[1]
(2, 5)
Note that the index starts at 0, so to get the second column as in your example you use zip(*array)[1] instead of zip(*array)[2]. zip() returns tuples instead of lists, depending on how you are using it this may not be a problem, but if you need lists you can always do map(list, zip(*array)) or list(zip(*array)[1]) to do the conversion.
If you use Matlab, you probably will want to install NumPy:
Using NumPy, you can do this:
In [172]: import numpy as np
In [173]: arr = np.matrix('1 2 3; 4 5 6')
In [174]: arr
Out[174]:
matrix([[1, 2, 3],
[4, 5, 6]])
In [175]: arr[:,2]
Out[175]:
matrix([[3],
[6]])
Since Python uses 0-based indexing (while Matlab uses 1-based indexing), to get the same slice you posted you would do:
In [176]: arr[:,1]
Out[176]:
matrix([[2],
[5]])
It is easy to build numpy arrays of higher dimension as well. You could use np.dstack for instance:
In [199]: B = np.dstack( (np.eye(3), np.ones((3,3)), np.arange(9).reshape(3,3)) )
In [200]: B.shape
Out[200]: (3, 3, 3)
In [201]: B[:,:,0]
Out[201]:
array([[ 1., 0., 0.],
[ 0., 1., 0.],
[ 0., 0., 1.]])
In [202]: B[:,:,1]
Out[202]:
array([[ 1., 1., 1.],
[ 1., 1., 1.],
[ 1., 1., 1.]])
In [203]: B[:,:,2]
Out[203]:
array([[ 0., 1., 2.],
[ 3., 4., 5.],
[ 6., 7., 8.]])
And here is the array formed from the second column from each of the 3 arrays above:
In [204]: B[:,1,:]
Out[204]:
array([[ 0., 1., 1.],
[ 1., 1., 4.],
[ 0., 1., 7.]])
Numpy doesn’t have a function to create magic squares, however. sniff