How to get the cartesian indexes of the k largest elements of a matrix (2D array) efficiently?

New to Julia
,
1

Hi! I would like to find a way to get the cartesian indexes of the k largest elements in a matrix (2D array). Do you have suggestions of a computationally efficient method to achieve that?

Example: for k=3 with the matrix A below, the cartesian indexes would be [(1, 3), (2, 2), (2, 3)].

julia> A = [1 1 3; 1 3 3]
2×3 Matrix{Int64}:
 1  1  3
 1  3  3

Thanks!

2

What Is the result expected for k=2?

3

I guess either ties are broken arbitrarily, or the first element is returned (I am interested in both solutions). By first element, I mean the first reached when going through the linear indexing of the matrix A (from 1 to 6 in this example).

4

Maybe something like

c = Vector{CartesianIndex{2}}(undef,k)
for i in lastindex(M):lastindex(M)-k
    c[i] = CartesianIndex...
end

(Sorry, on the cell phone…)

5

I am unsure I understood, but this made me think of 1) flatten the array, 2) sortperm to get the largest elements and 3) retrieve the cartesian indexes from the k last elements? This would give something like:

function klargest_indexes(m, k)
    ci = CartesianIndices(size(m))
    p = sortperm(vec(m))[end-k+1:end]
    ci[p]
end
k = 2
m = [1 2 3; 4 5 5; 5 4 3]
out = klargest_indexes(m, k)

In such a case, ties yield the last elements in the linear indexing. Is this what you hinted? Is it computationally efficient? I get a @btime of 5.374 μs (6 allocations: 3.55 KiB) for m = rand(20, 20), k = 2.

6

I meant probably something like this:

julia> ci = CartesianIndices(size(A))
       for i in lastindex(A)-2:lastindex(A) 
           @show ci[i]
       end
ci[i] = CartesianIndex(2, 2)
ci[i] = CartesianIndex(1, 3)
ci[i] = CartesianIndex(2, 3)

(you can by default run over a matrix as if it was a linearly indexed)

7

use instead

sortperm(vec(m), 1:k; rev=true)
8

I think you mean partialsortperm(vec(m), 1:k; rev=true)

4 Likes
9

Thanks! This is indeed more efficient!

10

Am I missing something, or this is much more efficient?

julia> A = [1 1 3; 1 3 3]
2×3 Matrix{Int64}:
 1  1  3
 1  3  3

julia> function getci(M,k)
           ci = CartesianIndices(size(M))
           cis = Vector{CartesianIndex{2}}(undef,k)
           j = 0
           for i in lastindex(M)-k+1:lastindex(M)
               j += 1
               cis[j] = ci[i]
           end
           return cis
       end
getci (generic function with 1 method)

julia> @btime getci($A,3)
  32.426 ns (1 allocation: 128 bytes)
3-element Vector{CartesianIndex{2}}:
 CartesianIndex(2, 2)
 CartesianIndex(1, 3)
 CartesianIndex(2, 3)

vs

julia> function klargest_indexes(m, k)
           ci = CartesianIndices(size(m))
           p = partialsortperm(vec(m), 1:k; rev=true)
           ci[p]
       end
klargest_indexes (generic function with 1 method)

julia> @btime klargest_indexes($A,3)
  329.935 ns (6 allocations: 384 bytes)
3-element Vector{CartesianIndex{2}}:
 CartesianIndex(2, 2)
 CartesianIndex(1, 3)
 CartesianIndex(2, 3)
11

I think the sorting is lacking in the proposed getci function. Correct me if I am wrong but I believe it to yield the cartesian indexes of the k last elements rather than the k largest elements.

1 Like
12

Ah, I see. In the example you posted those are the same, and I misunderstood your question. Thanks.

(I think you can do that without any sorting, just with an auxiliary array of size k, and running once over the elements of the matrix storing the indexes of the k larger)

2 Likes
13

There you go:

julia> function getci!(M,k,cmins,vmins) # inplace
           ci = CartesianIndices(size(M))
           for i in 1:k
               cmins[i] = ci[i]
               vmins[i] = M[i]
           end
           imin = findmin(vmins)[2]
           for i in firstindex(M)+k:lastindex(M)
               if M[i] > vmins[imin]
                   cmins[imin] = ci[i]
                   vmins[imin] = M[i]
                   imin = findmin(vmins)[2]
               end
           end
           return cmins, vmins
       end

       function getci(M,k) # allocating
           cmins = Vector{CartesianIndex{2}}(undef,k)
           vmins = Vector{eltype(M)}(undef,k)
           return getci!(M,k,cmins,vmins)
       end
getci (generic function with 1 method)

julia> A = [3 1 1
            1 3 3]
2×3 Matrix{Int64}:
 3  1  1
 1  3  3

julia> @btime getci($A,3)
  78.857 ns (2 allocations: 240 bytes)
(CartesianIndex{2}[CartesianIndex(1, 1), CartesianIndex(2, 2), CartesianIndex(2, 3)], [3, 3, 3])

julia> cmins=Vector{CartesianIndex{2}}(undef,3); vmins=Vector{Int}(undef,3);

julia> @btime getci!($A,3,$cmins,$vmins)
  35.864 ns (0 allocations: 0 bytes)
(CartesianIndex{2}[CartesianIndex(1, 1), CartesianIndex(2, 2), CartesianIndex(2, 3)], [3, 3, 3])

(since here there is an auxiliary array with the associated values, why not return it)

(for rand(20,20) and k=2, an example you mentioned this is 10 times faster than the alternative:

julia> A = rand(20,20); k = 2;

julia> @btime klargest_indexes($A,2)
  5.540 μs (8 allocations: 3.56 KiB)
2-element Vector{CartesianIndex{2}}:
 CartesianIndex(14, 8)
 CartesianIndex(10, 10)

julia> @btime getci($A,2)
  549.492 ns (2 allocations: 208 bytes)
(CartesianIndex{2}[CartesianIndex(10, 10), CartesianIndex(14, 8)], [0.9941985771967297, 0.9992607929792117])

julia> cmins=Vector{CartesianIndex{2}}(undef,k); vmins=Vector{eltype(A)}(undef,k);

julia> @btime getci!($A,2,$cmins,$vmins);
  509.870 ns (0 allocations: 0 bytes)
2 Likes