Search code examples
julia

How to split a vector according to a function in Julia?


In Julia, the findall(f,v) function gives me the collection of elements in a vector where f applied to the element is true. For example, findall(iseven,[1,2,3]) returns [2]. Now is there a function implemented in base (or a simple implementation) where I actually get the split vector instead of only the true instances? I mean, something like partition(iseven, [1,2,3]) returning ([2],[1,3]).

I can do a naive implementation, but I'm looking for a something efficient, preferably an existing function in Base...


Solution

  • One way could be:

    separate(p, x) = ( t = findall(p, x); 
      (@view(x[t]), @view(x[setdiff(eachindex(x), t)])) )
    

    and with separate you can do:

    julia> x = [1,2,3,5,6];
    
    julia> separate(iseven, x)
    ([2, 6], [1, 3, 5])
    
    julia> E, O = collect.(separate(iseven, x))
    ([2, 6], [1, 3, 5])
    
    julia> E
    2-element Vector{Int64}:
     2
     6
    

    Without collect. the function returns views which are good if copying the original vector is expensive.