【代码开发】选择一维数组中最大的k个

现在回顾本文，其实就是一个因为不熟悉库函数而自己去实现的愚蠢故事。

记住，下次自己想实现一些基础功能的时候，一定要先去看看能用上哪些库函数。

目标

在自己写代码的时候，碰到一个蛮有意思的函数。有点像算法题，目标是用torch和numpy的库函数实现。

先介绍一下函数的目标，就是从一个n行m列的数组中，选出每m个中最大的k个，返回一个n*k的数组，每个元素是对应k个元素的下标。返回下标可以乱序。

def select_maxk(arr, pick_num=1)
arr是一个n行m列的二维数组，pick_num就是上面提到的k的值，默认为1。

另外，还有epsilon = 0.1 的概率是随机选择k个元素的下标，来模拟一些特定的场景。（这需要在函数中随机生成一个向量来进行控制。）
比如生成向量tensor([False, True, False, False, False, False]) 就表示，除了第二行是随机返回k个下标，其余几行都必须返回最大元素的下标。

举个例子：n=6,m=4,k=3
原数组：
tensor([[0.1266, 0.1014, 0.1671, 0.6924],
[0.5884, 0.4361, 0.3644, 0.0175],
[0.8387, 0.7729, 0.0654, 0.1425],
[0.5068, 0.7436, 0.3154, 0.3503],
[0.4896, 0.6226, 0.1646, 0.6054],
[0.6791, 0.5005, 0.2655, 0.2428]])
每行最大的3个元素选出来后返回它们在本行的坐标：
tensor([False, True, False, False, False, False])
tensor([[0., 2., 3.],
[3., 1., 2.],
[3., 0., 1.],
[3., 0., 1.],
[0., 1., 3.],
[2., 1., 0.]])

即第一行最大的是第0、2、3个元素，第二行最大的是第0、1、2个元素（但因为对应的向量是True所以随机返回了1、2、3），依次类推。

如果k = 2则返回
tensor([False, False, False, False, False, False])
tensor([[2., 3.],
[1., 0.],
[0., 3.],
[0., 1.],
[3., 1.],
[1., 0.]])

函数实现

import numpy as np
import random
import torch as th
def select_maxk(arr, pick_num=1):
    row = arr.shape[0]
    line = arr.shape[1]

    epsilon = 0.1
    random_numbers = th.rand(row)
    pick_random = random_numbers < epsilon
    # print(pick_random)
    
    result = th.empty(row, pick_num)
    for i in range(row):
        if pick_random[i]:
            randk = th.tensor(random.sample(range(line), pick_num))
            result[i] = randk
        else:
            maxk = th.tensor(np.argpartition(arr[i].numpy().ravel(), -pick_num)[-pick_num:])
            result[i] = maxk
    return result

这是初稿。我这个实现不是最优的。最起码有这么几点可以改进：

1. 随机生成向量这件事可以丢给另外一个函数做
2. 不用for循环而采用矩阵的方式
3. 混用了numpy和torch，增加数据计算的负担

————————————————

优化改进

其实如果熟悉pytorch的小伙伴，知道topk这个函数的话，直接调库就可以搞定了。
https://pytorch.org/docs/stable/generated/torch.topk.html#torch.topk

torch.topk(input, k, dim=None, largest=True, sorted=True, *, out=None)

>>> x = torch.arange(1., 6.)
>>> x
tensor([ 1.,  2.,  3.,  4.,  5.])
>>> torch.topk(x, 3)
torch.return_types.topk(values=tensor([5., 4., 3.]), indices=tensor([4, 3, 2]))

这样的话我上面的算法就是一个伪命题了，不如说，如果一开始就知道这个函数，就没有这篇博客了。

一开始的时候我没有在网上找到类似的接口。结果就让自己费时又费力。
这个故事告诉我们，一定要熟悉库函数，没事多去看看库函数！

优化后函数

import numpy as np
import random
import torch as th
def select_topk(arr, pick_num=1):
    row = arr.shape[0]
    line = arr.shape[1]

    epsilon = 0.1
    random_numbers = th.rand(row)
    pick_random = random_numbers < epsilon

    _, result = th.topk(arr, pick_num, dim = 1)
    for i in range(row):
        if pick_random[i]:
            randk = th.tensor(random.sample(range(line), pick_num))
            result[i] = randk
    return result

总结

多去看看库函数！
多去看看库函数！
多去看看库函数！