C++ – 求解std::vector中topk数值以及topk数值对应的索引

1 C++获取vector中topk数值以及topk数值对应的索引

1.1 思路

topk问题是日常应用经常会遇到的问题，在python中对于多分类问题，经常需要输出top5的标签，在C++中，topk问题也是经常面对的问题。

比如给定一个std::vector<float>的容器，需要找出该容器中最大的top5的数值以及这5个数值在该容器中的索引，对于这个问题可以使用std::sort函数进行解决，我将求解TopK数值对应索引的函数进行了如下的封装：

std::vector<int> TopKIndex(const std::vector<float>& vec,int topk)
{
    std::vector<int> topKIndex;
    topKIndex.clear();

    std::vector<size_t> vec_index(vec.size());
    std::iota(vec_index.begin(), vec_index.end(), 0);

    std::sort(vec_index.begin(), vec_index.end(), [&vec](size_t index_1, size_t index_2) { return vec[index_1] > vec[index_2]; });

    int k_num = std::min<int>(vec.size(), topk);

    for (int i = 0; i < k_num; ++i)
    {
        topKIndex.emplace_back(vec_index[i]);
    }

    return topKIndex;
}

在上述代码中，先声明一个与原始数值容器等长的std::vector<size_t>容器，用于存储topk数值对应的索引，然后通过std::iota将该索引容器初始化为{0,1,2,......}的形式，然后使用std::sort方法通过比较原始数值容器的元素的大小去更新索引容器中索引的顺序。

1.2 示例

以下代码是求解std::vector<float>容器中top5数值以及对应索引的示例。

#include <iostream>
#include <algorithm>
#include <vector>
#include <numeric>
#include <iterator>

std::vector<int> TopKIndex(const std::vector<float>& vec,int topk)
{
    std::vector<int> topKIndex;
    topKIndex.clear();

    std::vector<size_t> vec_index(vec.size());
    std::iota(vec_index.begin(), vec_index.end(), 0);

    std::sort(vec_index.begin(), vec_index.end(), [&vec](size_t index_1, size_t index_2) { return vec[index_1] > vec[index_2]; });

    int k_num = std::min<int>(vec.size(), topk);

    for (int i = 0; i < k_num; ++i)
    {
        topKIndex.emplace_back(vec_index[i]);
    }

    return topKIndex;
}


int main()
{
    std::vector<float> testVec = { 1.0,0.9,2.0,89.6,3.0,4.0,4.1,4.3,115.3,4.9,5.9,10.2,11.6,9.5,25.3,99.3 };

    std::vector<int> indexResult = TopKIndex(testVec,5);

    std::cout << "原数组Top5数值为：" << std::endl;
    for (int i = 0; i < indexResult.size(); ++i)
    {
        std::cout << testVec[indexResult[i]] << std::endl;
    }

    std::cout << "原数组Top5数值对应的索引为：" << std::endl;
    for (int i = 0; i < indexResult.size(); ++i)
    {
        std::cout << indexResult[i] << std::endl;
    }

    int a = 1;

    return 0;
}

运行结果：