Element selection from array

Question

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

Look into the builtin zip

Brian61354270
– Brian61354270

2020-03-01 17:40:31 +00:00
Commented Mar 1, 2020 at 17:40 — Brian61354270
– Brian61354270, Commented Mar 1, 2020 at 17:40
Are you working with lists or np.arrays?

Tomerikoo
– Tomerikoo

2020-03-01 17:41:28 +00:00
Commented Mar 1, 2020 at 17:41 — Tomerikoo
– Tomerikoo, Commented Mar 1, 2020 at 17:41

Brian61354270 · Accepted Answer · 2020-03-01 17:50:13Z

3

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.

edited Mar 1, 2020 at 17:50

answered Mar 1, 2020 at 17:42

Brian61354270

14.9k7 gold badges37 silver badges64 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

user12938120 Over a year ago

Thanks this worked, and yeah the list is named differently in my code was just simplifying for on here

Tomerikoo · Accepted Answer · 2020-03-01 17:43:46Z

0

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

answered Mar 1, 2020 at 17:43

Tomerikoo

19.6k16 gold badges57 silver badges68 bronze badges

2 Comments

Boseong Choi Over a year ago

enumerate works, but I think using zip is more pythonic.

Tomerikoo Over a year ago

Yeah I agree that using zip saves the extra A[row]. I will leave this as another approach anyway...

Collectives™ on Stack Overflow

Element selection from array

2 Answers 2

1 Comment

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related