在Python中从数组中选择列

Question

In trying to select the first (when counting from zero) column in a 2D 4x4 array, I wrote the following script:

import numpy
a4x4=[list(range(4*i,4*(i+1))) for i in list(range(4))]
print(a4x4)
print(a4x4[:,1])

The array seems to be alright:

[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]

but instead of

[1, 5, 9, 13]

for the second print, I get this error:

TypeError: list indices must be integers, not tuple

Why does this error appear, what is going wrong?

Answer 1

You've import numpy but you aren't using it. What you have instead is a list of lists, and Python doesn't support multidimensional slicing for that (ie, you'd need [a4x4[i][1] for i in range(4)] to get the result you expect, but really you should be using numpy). Here's an example:

import numpy
a4x4=numpy.array([list(range(4*i,4*(i+1))) for i in list(range(4))])
print(a4x4)
print(a4x4[:,1])

By the way, in numpy you can also build the array you want directly, like this:

 numpy.arange(4*4).reshape((4,4))

(And also in Python one doesn't need the list calls I have above, I'm just trying to keep the code as similar to yours as possible to see the key thing, which is converting the list of lists into a numpy array.)

Answer 2

You can produce the result you want using list comprehension - just as you created the original 4x4:

a4x4=[list(range(4*i,4*(i+1))) for i in list(range(4))]
print([a4x4[i][1] for i in range(4)])

furthermore, you can simplify your logic a bit by tossing out the list function:

a4x4 = [range(4*i,4*(i+1)) for i in range(4)]
print([a4x4[i][1] for i in range(4)])