C++上二维向量的Argmax

Argmax of 2d vector on C++

本文关键字：二维 向量 Argmax C++      更新时间：2024-09-24

我正在研究python/pytorch，我有一个类似的例子

2d vector a

|
v
dim-0  ---> -----> dim-1 ------> -----> --------> dim-1
|   [[-1.7739,  0.8073,  0.0472, -0.4084],
v    [ 0.6378,  0.6575, -1.2970, -0.0625],
|    [ 1.7970, -1.3463,  0.9011, -0.8704],
v    [ 1.5639,  0.7123,  0.0385,  1.8410]]
|
v

然后，索引为1的argmax将是

# argmax (indices where max values are present) along dimension-1
In [215]: torch.argmax(a, dim=1)
Out[215]: tensor([1, 1, 0, 3])

我的问题是，给定如上所述的二维向量a，我如何在C++上实现argmax函数，以获得与上述相同的输出？感谢阅读这就是我所做的


vector<vector<float>> a_vect
{
{-1.7739,  0.8073,  0.0472, -0.4084},
{0.6378,  0.6575, -1.2970, -0.0625},
{1.7970, -1.3463,  0.9011, -0.8704},
{1.5639,  0.7123,  0.0385,  1.8410}
};
std::vector<int>::iterator max = max_element(a_vect.begin() , a_vect.end()-a_vect.begin());

您可以使用std::max_element在每个子向量中查找索引

#include <algorithm>
#include <iostream>
#include <vector>
using std::vector;
int main()
{
vector<vector<float>> a_vect=
{
{-1.7739,  0.8073,  0.0472, -0.4084},
{0.6378,  0.6575, -1.2970, -0.0625},
{1.7970, -1.3463,  0.9011, -0.8704},
{1.5639,  0.7123,  0.0385,  1.8410}
};

vector<int> max_index;
for(auto& v:a_vect)
max_index.push_back(std::max_element(v.begin(),v.end())-v.begin());

for(auto i:max_index)
std::cout << i << ' '; // 1 1 0 3 
}
相关文章：