Question

我有两个NumPy数组。在我的情况下，Y包含输出和P此输出正确的概率。行和列的形式（输出，noOfAnswers）或（概率，noOfAnswers）。所以一般来说输出比noOfAnswers大得多。

我正在通过以下方式选择有关P的两个最重要的结果：

chooseThem = np.argpartition(P,-2,axis=1)[:,-2:]

现在我希望创建一个大小（输出，2）的新数组YP，只包含chooseThem指定的值。使用for循环，这很简单，但性能不佳。

这是一个使用一些人工数组的“坏方法”的例子：

import numpy as np
Y = 4*(np.random.rand(1000,6)-0.5)
P = np.random.rand(1000,6)
biggest2 = np.argpartition(P,-2,axis=1)[:,-2:]
YNew = np.zeros((1000,2))

for j in range(2):
    for i in range(1000):
        YNew[i,j] = Y[i,biggest2[i,j]]

有没有人建议快速创建这个新阵列？

Answer 1

这适用于切片数组

dex = np.array([np.arange(1000),np.arange(1000)]).T
YNew = Y[dex,biggest2]

进行一些测试（old = loop method new = index method）

1000行

%timeit new(Y,P,1000,biggest2)
The slowest run took 4.47 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 39.1 µs per loop

%timeit old(Y,P,1000,biggest2)
1000 loops, best of 3: 853 µs per loop

100000行

%timeit new(Y,P,100000,biggest2)
100 loops, best of 3: 4.49 ms per loop

%timeit old(Y,P,100000,biggest2)
10 loops, best of 3: 89.4 ms per loop

根据另一个numpy数组选择numpy数组中的行/列（性能）

1 个答案: