numpy - выбрать несколько элементов из каждой строки массива

Question

Мне нужно выбрать несколько разных значений в каждой строке 2D-массива.

A = np.array([[ 1, 2, 3, 4],
              [ 5, 6, 7, 8],
              [ 9,10,11,12])
A[something]

>>> np.array([[ 1, 2],
              [ 6, 7],
              [11,12]])

Я знаю, что могу создать логический массив такой же формы, как A, и установить каждый элемент в цикле for, ноЯ надеюсь найти лучшее решение.

Answer 1

Вы можете попробовать следующее:

import numpy as np
A = np.array([[ 1, 2, 3, 4],
              [ 5, 6, 7, 8],
              [ 9,10,11,12]])
i = [[0],[1],[2]]
j = [[0,1], [1,2],[2,3]]
B = A[i,j]
print(B)
#Prints
[[ 1  2]
 [ 6  7]
 [11 12]]