python - 返回 numpy 数组的最小 X 数并保持顺序

Question

我有以下 X numpy 数组。我想从 X[i] 创建一个数组，其中包含数组的三个最小 X[i][3] 值。

array([[  2, 356,   1, 0.7],
       [  3, 356,   1, 5],
       [  3, 357,   1, 3],
       [  4, 355,   1, 0.1],
       [  4, 356,   1, 16],
       [  4, 357,   1, 2]])

结果应该是这样的:

array([[  2, 356,   1, 0.7],
       [  4, 355,   1, 0.1],
       [  4, 357,   1, 2]])

Answer 1

这是一种方法-

X[np.sort(X[:,3].argsort()[:3])]

基本上，我们使用argsort 来获取排序的索引，选择前三个作为最低的三个元素。我们将使用这些索引为输出索引数组。要保持输入数组中的顺序，请在索引之前对这些索引进行排序。

sample 运行-

In [148]: X
Out[148]: 
array([[  2.00e+00,   3.56e+02,   1.00e+00,   7.00e-01],
       [  3.00e+00,   3.56e+02,   1.00e+00,   5.00e+00],
       [  3.00e+00,   3.57e+02,   1.00e+00,   3.00e+00],
       [  4.00e+00,   3.55e+02,   1.00e+00,   1.00e-01],
       [  4.00e+00,   3.56e+02,   1.00e+00,   1.60e+01],
       [  4.00e+00,   3.57e+02,   1.00e+00,   2.00e+00]])

In [149]: X[np.sort(X[:,3].argsort()[:3])]
Out[149]: 
array([[  2.00e+00,   3.56e+02,   1.00e+00,   7.00e-01],
       [  4.00e+00,   3.55e+02,   1.00e+00,   1.00e-01],
       [  4.00e+00,   3.57e+02,   1.00e+00,   2.00e+00]])

为了性能，我们可以使用np.argpartition .因此，X[:,3].argsort()[:3] 可以替换为 np.argpartition(X[:,3],3)[:3]。 argpartition 因为它的实现方式为我们提供了对应于最低 3 元素的索引，只是不一定按照最低到第二低到第三低的顺序。但这没关系，因为无论如何我们稍后都会对这些索引进行排序，以保持输入数组中的顺序(前面讨论过)。

性能提升建议的时间 -

In [164]: X = np.random.rand(100000,4)

In [165]: np.sort(X[:,3].argsort()[:3])
Out[165]: array([ 9950, 69008, 76552])

In [166]: np.sort(np.argpartition(X[:,3],3)[:3])
Out[166]: array([ 9950, 69008, 76552])

In [167]: %timeit np.sort(X[:,3].argsort()[:3])
100 loops, best of 3: 7.59 ms per loop

In [168]: %timeit np.sort(np.argpartition(X[:,3],3)[:3])
1000 loops, best of 3: 290 µs per loop