Find n smallest values in an n dims array

New to Julia
, , , ,
21 
julia> A = rand(0x0:0x8f, 1_000_000);

julia> @which findmin(A)
findmin(A::AbstractArray; dims) in Base at reducedim.jl:1005

julia> @less findmin(A)

findmin(A::AbstractArray; dims=:) = _findmin(A, dims)

function _findmin(A, region)
    ri = reduced_indices0(A, region)
    if isempty(A)
        if prod(map(length, reduced_indices(A, region))) != 0
            throw(ArgumentError("collection slices must be non-empty"))
        end
        (similar(A, ri), zeros(eltype(keys(A)), ri))
    else
        findminmax!(isgreater, fill!(similar(A, ri), first(A)),
                    zeros(eltype(keys(A)), ri), A)
    end
end

The code:

22

Find them as below, for example the first 3 values.

julia> a = [100 1 10; 20 200 2]
2×3 Matrix{Int64}:
 100    1  10
  20  200   2

julia> first(a |> vec |> unique |> sort, 3)
3-element Vector{Int64}:
  1
  2
 10
1 Like
23

Hi !
I’ve used your code on my work, so I am submitting the slightly improved version I’m using:

function argsmallest(A::AbstractArray{T,N}, n::Integer) where {T,N}
    # should someone ask more elements than array size, just sort array
    if n>= length(vec(A))
        ind=collect(1:length(vec(A)))
        ind=sortperm(A[ind])
        return CartesianIndices(A)[ind]
    end

    # otherwise 
    ind=collect(1:n)
    mymax,getout=findmax(@view A[1:n])
    for j=n+1:length(vec(A))
        if A[j]<mymax
            ind[getout]=j
            mymax,getout=findmax(@view A[ind])
        end
    end
    
    return ind
end

I just add a tracker for mymax index and some @view to avoid memory allocation on access

3 Likes