Split vector into N (potentially unequal length) subvectors?

Is there a convenience function for splitting a vector into N subvectors? For example, if you had a length-10 vector and called the function to split it into 3 parts, you would get three vectors back, two with length 3 and another with length 4. The subvectors contain all the original items. I see many references to IterTools.partition but it performs a slightly different task.

I could write a little function to do this, but I wanted to see if one is already available.

Iterators.partition is the only builtin I’m aware of that comes close to this purpose. To get the “long tail” partition, rather than the “short tail” given by Iterators.partition, I’d do something like this:

function partitionvec(x,stride,longtail::Bool=true)
    # longtail=true to lengthen the last entry with the leftovers
    # longtail=false to place the leftovers in their own entry
    stride > 0 || error("stride must be positive") # doesn't handle negative strides
    starts = firstindex(x):stride:lastindex(x)-longtail*stride # where to start each subvector
    return [view(x,starts[i]:get(starts,i+1,lastindex(x)+1)-1) for i in eachindex(starts)]
end

Yeah I decided just to write a really explicit version

function makechunks(X::AbstractVector{T}, n::Int) where {T}
    L = length(X)
    c = L ÷ n
    Y = Vector{Vector{T}}(undef, n)
    idx = 1
    for i ∈ 1:n-1
        Y[i] = X[idx:idx+c-1]
        idx += c 
    end
    Y[end] = X[idx:end]
    return Y
end

You might want to use views instead to avoid copies. Also, the whole thing can be made more compact with a comprehension:

@views function makechunks(X::AbstractVector, n::Integer)
    c = length(X) ÷ n
    return [X[1+c*k:(k == n-1 ? end : c*k+c)] for k = 0:n-1]
end
3 Likes

Yes that’s all fair. It’s not a performance-critical bit of code for me at the moment so I wasn’t bothering. I’m just using it to split up some items before calling a function on the groups within a @threads loop.