np.partition 分区函数

np.partition 是对向量、沿着一个维度方向、按照大小对数据进行分堆，分成了两堆。

• 在 kth 前面的这堆数值，都是这个向量里面比较小的群体们。

• 在 kth 后面的这堆数值，都是这个向量里面比较大的群体们。

下面以示例展示：

import numpy as np

# 创建数据
a = np.array([[2, 3, 6, 2, 5, 1], 
              [7, 1, 76, 333, 3, 1], 
              [2, 3, 1, 2, 0, 1]])

沿着 axis=0 进行排序：

np.partition(a, 2, axis=0)

array([[  2,   1,   1,   2,   0,   1],
       [  2,   3,   6,   2,   3,   1],
       [  7,   3,  76, 333,   5,   1]])

沿着 axis=1 排序：

np.partition(a, 4, axis=1)

array([[  2,   1,   2,   3,   5,   6],
       [  3,   1,   1,   7,  76, 333],
       [  0,   1,   1,   2,   2,   3]])

也可以取出 topk 数据：

k = 3
np.partition(a, -k, axis=1)[:, -k:]

array([[  3,   5,   6],
       [  7, 333,  76],
       [  2,   2,   3]])