Julia idiomatic way to split vector to subvectors based on condition

user_5224881 0
Viewed 795

Let's say I have a vector a = [1, 0, 1, 2, 3, 4, 5, 0, 5, 6, 7, 8, 0, 9, 0] and I want to split it to smaller vectors based on a condition depending on value in that array. E.g. value being zero. Thus I want to obtain vector of following vectors

 [1, 0]
 [1, 2, 3, 4, 5, 0]
 [5, 6, 7, 8, 0]
 [9, 0]

So far this was working for me as a naive solution, but it loses the type.

function split_by_λ(a::Vector, λ)
    b = []
    temp = []
    for i in a
        push!(temp, i)
        if λ(i)
            push!(b, temp)
            temp = []
        end
    end
    b
end
split_by_λ(a, isequal(0))

Then I tried to play with ranges, which feels a bit more idiomatic, and does not lose the type.

function split_by_λ(a::Vector, λ)
    idx = findall(λ, a)
    ranges = [(:)(i==1 ? 1 : idx[i-1]+1, idx[i]) for i in eachindex(idx)]
    map(x->a[x], ranges)
end

split_by_λ(a, isequal(0))

but it still feels very cumbersome regarding it's a rather simple task. Is there something I'm missing, some easier way?

julia idioms
2 Answers

Maybe someone has a shorter idea but here is mine:

julia> inds = vcat(0,findall(==(0),a),length(a))


julia> getindex.(Ref(a), (:).(inds[1:end-1].+1,inds[2:end]))
5-element Array{Array{Int64,1},1}:
 [1, 0]
 [1, 2, 3, 4, 5, 0]
 [5, 6, 7, 8, 0]
 [9, 0]
 []

Or if you want to avoid copying a

julia> view.(Ref(a), (:).(inds[1:end-1].+1,inds[2:end]))
5-element Array{SubArray{Int64,1,Array{Int64,1},Tuple{UnitRange{Int64}},true},1}:
 [1, 0]
 [1, 2, 3, 4, 5, 0]
 [5, 6, 7, 8, 0]
 [9, 0]
 0-element view(::Array{Int64,1}, 16:15) with eltype Int64
user_99577100

Pretty much the same as Przemyslaw's answer, but maybe less cryptic dense:

function split_by(λ, a::Vector)
    first, last = firstindex(a), lastindex(a)
    splits = [first-1; findall(λ, a); last]
    s1, s2 = @view(splits[1:end-1]), @view(splits[2:end])
    return [view(a, i1+1:i2) for (i1, i2) in zip(s1, s2)]
end

Also, I changed the signature to the conventional one of "functions first", which allows you to use do-blocks. Additionally, this should work with offset indexing.

One could surely get rid of the intermediate allocations, but I think that gets ugly without yield:

function split_by(λ, a::Vector)
    result = Vector{typeof(view(a, 1:0))}()
    l = firstindex(a)
    r = firstindex(a)
    while r <= lastindex(a)
        if λ(a[r])
            push!(result, @view(a[l:r]))
            l = r + 1
        end
        r += 1
    end
    push!(result, @view(a[l:end]))

    return result
end
user_13462760
Related