Как получить доступ к i-му столбцу многомерного массива NumPy?

Предположим, у меня есть:

test = numpy.array([[1, 2], [3, 4], [5, 6]])

test[i]возвращает мне i-ю строку массива (например [1, 2]). Как я могу получить доступ к i-му столбцу? (например [1, 3, 5]). Кроме того, это будет дорогая операция?

>>> test[:,0]
array([1, 3, 5])

По аналогии,

>>> test[1,:]
array([3, 4])

позволяет вам получить доступ к строкам. Это описано в Разделе 1.4 (Индексирование) ссылки NumPy . Это быстро, по крайней мере, по моему опыту. Это, конечно, намного быстрее, чем доступ к каждому элементу в цикле.

И если вы хотите получить доступ к нескольким столбцам одновременно, вы можете сделать:

>>> test = np.arange(9).reshape((3,3))
>>> test
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])
>>> test[:,[0,2]]
array([[0, 2],
       [3, 5],
       [6, 8]])

>>> test[:,0]
array([1, 3, 5])

эта команда дает вам вектор строки, если вы просто хотите зациклить его, это нормально, но если вы хотите выполнить hstack с другим массивом с размером 3xN, у вас будет

ValueError: all the input arrays must have same number of dimensions

пока

>>> test[:,[0]]
array([[1],
       [3],
       [5]])

дает вам вектор-столбец, так что вы можете выполнять конкатенацию или операцию hstack.

например

>>> np.hstack((test, test[:,[0]]))
array([[1, 2, 1],
       [3, 4, 3],
       [5, 6, 5]])

Вы также можете транспонировать и возвращать строку:

In [4]: test.T[0]
Out[4]: array([1, 3, 5])

Чтобы получить несколько независимых столбцов, просто:

> test[:,[0,2]]

вы получите столбцы 0 и 2

Хотя на вопрос дан ответ, позвольте мне упомянуть некоторые нюансы.

Допустим, вы заинтересованы в первом столбце массива

arr = numpy.array([[1, 2],
                   [3, 4],
                   [5, 6]])

Как вы уже знаете из других ответов, чтобы получить его в виде «вектора строки» (массива формы (3,)), вы используете нарезку:

arr_c1_ref = arr[:, 1]  # creates a reference to the 1st column of the arr
arr_c1_copy = arr[:, 1].copy()  # creates a copy of the 1st column of the arr

Чтобы проверить, является ли массив представлением или копией другого массива, вы можете сделать следующее:

arr_c1_ref.base is arr  # True
arr_c1_copy.base is arr  # False

см. ndarray.base .

Помимо очевидного различия между ними (изменение arr_c1_refповлияет arr), количество байтовых шагов для прохождения каждого из них различно:

arr_c1_ref.strides[0]  # 8 bytes
arr_c1_copy.strides[0]  # 4 bytes

видеть успехи . Почему это важно? Представьте, что у вас есть очень большой массив Aвместо arr:

A = np.random.randint(2, size=(10000,10000), dtype='int32')
A_c1_ref = A[:, 1] 
A_c1_copy = A[:, 1].copy()

и вы хотите вычислить сумму всех элементов первого столбца, т.е. A_c1_ref.sum()A_c1_copy.sum() . или . Использование скопированной версии намного быстрее:

%timeit A_c1_ref.sum()  # ~248 µs
%timeit A_c1_copy.sum()  # ~12.8 µs

Это связано с разным количеством шагов, упомянутых ранее:

A_c1_ref.strides[0]  # 40000 bytes
A_c1_copy.strides[0]  # 4 bytes

Хотя может показаться, что лучше использовать копии столбцов, это не всегда верно по той причине, что создание копии занимает много времени и требует больше памяти (в этом случае мне понадобилось около 200 мкс для создания A_c1_copy). Однако, если нам нужна копия в первую очередь, или нам нужно выполнить много разных операций с определенным столбцом массива, и мы в порядке с потерей памяти ради скорости, то создание копии - это путь.

В случае, если мы заинтересованы работать в основном со столбцами, было бы неплохо создать наш массив в мажорном столбце ('F'), а не в мажорном ряду ('C') (по умолчанию ), а затем выполните нарезку, как раньше, чтобы получить столбец, не копируя его:

A = np.asfortranarray(A)  # or np.array(A, order='F')
A_c1_ref = A[:, 1]
A_c1_ref.strides[0]  # 4 bytes
%timeit A_c1_ref.sum()  # ~12.6 µs vs ~248 µs

Теперь выполнение операции суммирования (или любой другой) в представлении столбцов выполняется намного быстрее.

Наконец, позвольте мне заметить, что транспонирование массива и использование нарезки строк - это то же самое, что и нарезка столбцов в исходном массиве, потому что транспонирование выполняется просто путем замены формы и шагов исходного массива.

A.T[1,:].strides[0]  # 40000

>>> test
array([[0, 1, 2, 3, 4],
       [5, 6, 7, 8, 9]])

>>> ncol = test.shape[1]
>>> ncol
5L

Затем вы можете выбрать 2-4 столбец следующим образом:

>>> test[0:, 1:(ncol - 1)]
array([[1, 2, 3],
       [6, 7, 8]])