首页 > 解决方案 > Julia using argmax to return index within a row of a matrix instead of a CartesianIndex

问题描述

Say I have a matrix

X = [1 2 3 4; 1 4 3 2];

I would like to find the argmax of each row of this matrix, relative to that row and not the index of the entry within the entirety of X. Meaning that I want the output of argmax(X, dims = (2)) to be a vector,

[4, 2];

but the current output is an array of CartesianIndex

[CartesianIndex(1, 4), CartesianIndex(2, 2)];

Is there a way to specify this in the argmax function or to transform the output efficiently to my desired state?

标签: julia

解决方案


您可以使用eachrow迭代矩阵的行:

julia> argmax.(eachrow(X))
2-element Vector{Int64}:
 4
 2

推荐阅读