Extracting indexes of vectors in Julia

How to obtain the indexes of the nonzero terms of a vector in Julia?

x = [[3, 0, 0], [0, 4, 0], [5, 6, 0]]
output: ([0, 1, 2, 2]), [0, 1, 0, 1])

Standard Julia Vectors/Arrays use indexing that starts at 1.

If all of your subvectors have the same length, you may prefer to work with a matrix:

M = vcat(x'...)
idx = M |> x->findall(!=(0),x)

Otherwise, you can write a loop:

function findNonzero(x)
       I,J = Int[],Int[]
       for (i,v) in enumerate(x)
           for (j,w) in enumerate(v)
               if w!=0
                   push!(I,i); push!(J,j)
               end
           end
       end
       return I,J
end
1 Like

You can get that exact output from

function f(x)
    a = [(i - 1, j - 1) for (i, r) in enumerate(x) for j in findall(!=(0), r)]
    return first.(a), last.(a)
end

but there’s a lot that’s unidiomatic about wanting that output so you should probably take a step back and look at a bigger picture of what you’re trying to accomplish.

I would just write that as idx = findall(!=(0), M).

3 Likes

Another option that works for subvectors of arbitrary lengths (kept Julia convention of indices starting from 1):

function g(x)
    a = Int64[];  b = similar(a)
    for (i,xi) in pairs(x)
        idx = findall(!=(0), xi)
        !isempty(idx) && (for j in idx; push!(a,i); push!(b,j) end)
    end
    return a, b
end

You can use reduce(vcat,...) for all tuples of nonzeros and then get first and last of them.

A = reduce(vcat, tuple.(i,findall(!=(0),j)) for (i,j) in pairs(x))
4-element Vector{Tuple{Int64, Int64}}:
 (1, 1)
 (2, 2)
 (3, 1)
 (3, 2)

first.(A)
4-element Vector{Int64}:
 1
 2
 3
 3

last.(A)
4-element Vector{Int64}:
 1
 2
 1
 2