Element selection from array

user_12989322 0
Viewed 38

I have a list 

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

and A is the array:

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

I basically want to create an array which appends the element of each row according to the index position in the list index. So to create the array [5, 6, 7, 3, 4]. I tried the following nested for loop but it obviously returned the 5 values for each row, rather than the specific one from each row.

list = []
for i in index:
    for a in A:
         list.append(a[i])
matrix arrays python
2 Answers

Use the builtin zip to iterate both index and A simultaneously:

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

For added brevity, I've expressed your loop as a list comprehension. Manually iterating over the iterator produced by zip is also perfectly valid:

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

Also note that it is generally a poor practice to create a variable with same name as a builtin, so I'd advise you rename list to e.g. result, my_list, or a more contextually relevant name.

user_110821650

The index of elements in index is the row and their value is the column, so using enumerate you can do:

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

Or as list-comp:

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