Выбор элемента из массива

Question

У меня есть список

index = [1, 1, 1, 0, 0] 

, а A - это массив:

A = [[ 0.  5. 10. 15.]
     [ 1.  6. 11. 16.]
     [ 2.  7. 12. 17.]
     [ 3.  8. 13. 18.]
     [ 4.  9. 14. 19.]]

Я хочу создать массив, который добавляет элемент каждой строки в соответствии с позицией индекса в списке индексов. Итак, для создания массива [5, 6, 7, 3, 4]. Я попробовал следующее вложенное для l oop, но оно, очевидно, вернуло 5 значений для каждой строки, а не конкретное c по одному в каждой строке.

list = []
for i in index:
    for a in A:
         list.append(a[i])

Answer 1

Используйте встроенный zip для итерации одновременно и index и A:

result = [row[i] for i, row in zip(index, A)]

Для большей краткости я выразил ваш l oop как список понимания . Ручная итерация по итератору, производимому zip, также вполне допустима:

result = []
for i, row in zip(index, A):
    result.append(row[i])

Также обратите внимание, что обычно создание переменной с таким же именем, как у встроенного, плохая практика, поэтому я бы посоветовал вам переименуйте list, например, result, my_list или более релевантное для контекста имя.

Answer 2

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

res = []
for row, col in enumerate(index):
    res.append(A[row][col])

Или как list-comp :

res = [A[row][col] for row, col in enumerate(index)]