Wybierz jeden element w każdym wierszu tablicy Numpy według indeksów kolumn

Question

Czy istnieje lepszy sposób na uzyskanie "output_array" z "input_array" i "select_id"?

Czy możemy się pozbyć range(input_array.shape[0])?

>>> input_array = numpy.array([ [3,14], [12, 5], [75, 50] ]) 
>>> select_id = [0, 1, 1] 
>>> print input_array 
[[ 3 14] 
[12 5] 
[75 50]] 

>>> output_array = input_array[ range(input_array.shape[0]), select_id ] 
>>> print output_array 
[ 3 5 50] 

Answer 1

Można wybrać z danej tablicy przy użyciu numpy.choose który konstruuje tablicę z tablicy indeksu (w przypadku select_id) oraz zestaw tablic (w przypadku input_array) do wyboru. Jednak możesz najpierw przetłumaczyć input_array, aby dopasować wymiary. Poniżej przedstawiono mały przykład:

In [101]: input_array 
Out[101]: 
array([[ 3, 14], 
     [12, 5], 
     [75, 50]]) 

In [102]: input_array.shape 
Out[102]: (3, 2) 

In [103]: select_id 
Out[103]: [0, 1, 1] 

In [104]: output_array = np.choose(select_id, input_array.T) 

In [105]: output_array 
Out[105]: array([ 3, 5, 50]) 

Answer 2

Jak o:

[input_array[x,y] for x,y in zip(range(len(input_array[:,0])),select_id)] 

Answer 3

myślę enumerate jest poręczny.

[input_array[enum, item] for enum, item in enumerate(select_id)] 

Answer 4

(bo nie mogę pisać to jako komentarz na temat przyjętego odpowiedź)

Zauważ, że numpy.choose działa tylko jeśli masz 32 lub mniej wyborów (w tym przypadku wymiar swoje tablica, w której indeksujesz, musi mieć rozmiar 32 lub mniejszy). Dodatkowo, documentation dla numpy.choose mówi, że "Aby zmniejszyć prawdopodobieństwo błędnej interpretacji, nawet jeśli następujące" nadużycie "jest nominalnie wspierane, wybory nie powinny być ani nie powinny być uważane za pojedynczą tablicę, tj. Najbardziej zewnętrzną sekwencję. podobnie jak kontener powinien być listą lub krotką. "

PO pyta:

1. Czy istnieje lepszy sposób, aby uzyskać "output_array" z "input_array" i "select_id"?
   • Powiedziałbym, że sposób, który pierwotnie sugerowałeś, wydaje się najlepszy z prezentowanych tutaj. Jest łatwy do zrozumienia, skaluje się do dużych tablic i jest wydajny.
2. Czy możemy pozbyć się zakresu (input_array.shape [0])?
   • Tak, jak pokazują inne odpowiedzi, ale zaakceptowana nie działa tak ogólnie, jak sugeruje to już PO.