numpy如何從每一行中取一個元素

Question

假設我有形狀(n,m) numpy數組A 我想通過從A每一行中選擇一個元素來制作形狀為(n,)的數組b 。 我制作了一個形狀為(n,)的索引y的向量，我想構造b以便每個i b[i] = A[i,y[i]] 。

我可以想到的一種方法是b = A[np.arange(n), y] ，它b = A[np.arange(n), y]我的要求。 但是，我想知道是否有更有效的方法來執行此操作，因為我正在用不同的A多次調用該函數的函數中調用此操作，並且不希望python構造np.arange(n)每次調用該函數。

Answer 1

這是一個簡化的答案，但是如果您可以保持np.arange左右的速度會更快，但是可能更快的方法是計算y以匹配所需的A的扁平索引（即y + = np.arange（ N）×M）。 以下代碼對每個解決方案進行計時，並給出相同的輸出。 *Note: If your arrays are too large then flatten will take some time and keeping arange() array around will be faster

from timeit import timeit

setup = '''
import numpy as np
k = 6
y = np.arange(k)
a = np.arange(k**2).reshape(k,k)
'''

setup2 = '''import numpy as np
k = 6
y = np.arange(k)
a = np.arange(k**2).reshape(k,k)
y += y * k
'''

print "Original:", timeit("b = a[np.arange(k),y]",setup=setup)
print "Keep np.arange around:", timeit("b = a[y,y]",setup=setup)
print "Flatten:", timeit("b = a.flatten()[y]", setup=setup2)
#Original: 1.57689285278
#Keep np.arange around: 0.977301120758
#Flatten: 0.710500001907