Question

我正在使用python 3.6和numpy。我有一个n维数组。我需要在数组的最后一个维度上执行partition和argpartition。我显然可以调用这两个函数，但感觉就像浪费资源一样。有没有办法同时获得np.partition和np.argpartition的结果？应该有办法让np.partition的结果应用于数组我从np.argpartition获得的索引，但我现在还没有看到它！

谢谢！

Answer 1

获取这些argpartition索引，然后使用advanced-indexing获取分区数组。

因此，任何维度和任何通用轴的通用ndarray的实现都是这样的 -

def partition_results(a, k, axis=-1):
    idx = np.argpartition(a, k, axis=axis)
    index_arr = list(np.ix_(*[range(i) for i in a.shape]))
    index_arr[axis] =  idx
    return idx, a[index_arr]

np.ix_为我们提供了＆＃34;展开＆＃34;范围数组来完成advanced-indexing的任务。需要这些范围数组来覆盖与argpartition索引数组中除了最后一个之外的轴的长度相对应的所有维度，我们自己拥有这些argpartition索引。这种索引操作需要这种设置。

所以，通过使用对np.argpartition和np.partition的两个单独调用的方法，我们可以这样做 -

def partition_results_exclusive_way(a, k):
    idx = np.argpartition(a, k, axis=-1)
    part_arr = np.partition(a, k, axis=-1)
    return idx , part_arr

我们将在下一节中将其用于性能和价值验证的比较。

样本运行和运行时测试 -

In [496]: a = np.random.rand(20,20,20,20,20)

In [502]: A0, B0 = partition_results_exclusive_way(a, 10)

In [503]: A1, B1 = partition_results(a, 10)

In [504]: np.allclose(A0,A1)
Out[504]: True

In [505]: np.allclose(B0,B1)
Out[505]: True

In [506]: %timeit partition_results_exclusive_way(a, 10)
10 loops, best of 3: 92.6 ms per loop

In [507]: %timeit partition_results(a, 10)
10 loops, best of 3: 76 ms per loop

详细说明性能数字，让我们分开时间argpartition和partition -

In [509]: %timeit np.argpartition(a, 10, axis=-1)
10 loops, best of 3: 49.6 ms per loop

In [510]: %timeit np.partition(a, 10, axis=-1)
10 loops, best of 3: 43.6 ms per loop

因此，advanced-indexing操作花费了我们大约一半np.partition的费用。我们肯定在那里拯救！

同时获得numpy.partition和numpy.argpartition输出

1 个答案: