Як отримати доступ до 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])

ця команда дає вам вектор рядка, якщо ви просто хочете перевести цикл на нього, це добре, але якщо ви хочете встановити інший масив із розмірністю 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') замість порядку-рядка-головного ("С") (що є типовим) ), а потім виконайте нарізку, як раніше, щоб отримати стовпець, не копіюючи його:

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]])