Search code examples
performancejuliaargmax

Find index of maximum element satisfying condition (Julia)

In Julia I can use argmax(X) to find max element. If I want to find all element satisfying condition C I can use findall(C,X). But how can I combine the two? What's the most efficient/idiomatic/concise way to find maximum element index satisfying some condition in Julia?

Solution

  • If you'd like to avoid allocations, filtering the array lazily would work:

    idx_filtered = (i for (i, el) in pairs(X) if C(el))
argmax(i -> X[i], idx_filtered)

    Unfortunately, this is about twice as slow as a hand-written version. (edit: in my benchmarks, it's 2x slower on Intel Xeon Platinum but nearly equal on Apple M1)

    function byhand(C, X)
    start = findfirst(C, X)
    isnothing(start) && return nothing

    imax, max = start, X[start]
    for i = start:lastindex(X)
        if C(X[i]) && X[i] > max
            imax, max = i, X[i]
        end
     end
     imax, max
end