Python/Numpy: Selecting a Specific Column in a 2D Array
Join the DZone community and get the full member experience.Join For Free
I’ve been playing around with numpy this evening in an attempt to improve the performance of a Travelling Salesman Problem implementation and I wanted to get every value in a specific column of a 2D array.
The array looked something like this:
>> x = arange(20).reshape(4,5) >>> x array([[ 0, 1, 2, 3, 4], [ 5, 6, 7, 8, 9], [10, 11, 12, 13, 14], [15, 16, 17, 18, 19]])
I wanted to get the values for the 2nd column of each row which would return an array containing 1, 6, 11 and 16.
For some reason I was expecting it to be quite complicated but in fact we can do this by using matrix style syntax like so:
>>> x[:, 1] array([ 1, 6, 11, 16])
Here we are first saying that we want to return all the rows by specifying ‘:’ and then the ’1′ indicates that we only want to return the column with index 1.
If we wanted to return a specific row as well then we’d specify a value before the comma and it’d be a standard 2D array value lookup:
>> x[2,1] 11
>> x 11
Published at DZone with permission of Mark Needham, DZone MVB. See the original article here.
Opinions expressed by DZone contributors are their own.