Get the indices and corresponding combinations that contain given term from a list of combinations with replacement

Question

Let's say I have an ordered list of combinations with replacement with n=4 (objects) and r= 3 (sample):

[1, 1, 1]
 [1, 1, 2]
 [1, 1, 3]
 [1, 1, 4]
 [1, 2, 2]
 [1, 2, 3]
 [1, 2, 4]
 [1, 3, 3]
 [1, 3, 4]
 [1, 4, 4]
 [2, 2, 2]
 [2, 2, 3]
 [2, 2, 4]
 [2, 3, 3]
 [2, 3, 4]
 [2, 4, 4]
 [3, 3, 3]
 [3, 3, 4]
 [3, 4, 4]
 [4, 4, 4]

How could I get the indices of all the lists that contain a given object and generate all those corresponding lists (order of the elements inside the generated lists doesn't matter but the lists must be in the same order as the indices).

E.g.

Given element 2 and the array above I would want the function to return

• indices (in 1-based indexing): 2, 5,6,7,11,12,13,14,15,16
• an ordered list of lists/vectors (though the elements inside the lists/vectors don't need to be in this order):

[[1, 1, 2], 
[1, 2, 2],
[1, 2, 3]
[1, 2, 4]
[2, 2, 2]
[2, 2, 3]
[2, 2, 4]
[2, 3, 3]
[2, 3, 4]
[2, 4, 4]]

Of course I could just create this list every time and choose the required elements etc but I was wondering if there is a formula to do this - saving memory etc in case of very large arrays.

Answer 1

Here's one approach. It recognizes that you can calculate the next item in the list of combinations purely from the current item. To that end, it uses a next! function that mutates the current item into the next one.

I suppose this could be made cleaner by implementing an iteration interface, but this should work:

function next!(v, N, n=lastindex(v))
    val = v[n]
    if val < N
        v[n:end] .= (val + 1)
        status = true
    elseif n > firstindex(v)
        (v, status) = next!(v, N, n-1)
    else
        status = false
    end
    return (v, status)
end

function get_from_ordered_list((N, n), val)
    v, status = ones(Int, n), true
    i = 1
    inds = Int[]
    vecs = Vector{Int}[]
    while status
        if val in v
            push!(inds, i)
            push!(vecs, copy(v))
        end
        (v, status) = next!(v, N)
        i += 1
    end
    return (inds, vecs)
end

You could improve on this, performancewise by working on statically sized containers, but I think this should be ok. Benchmark:

julia> @btime get_from_ordered_list((4, 3), 2)
  666.883 ns (17 allocations: 1.78 KiB)
([2, 5, 6, 7, 11, 12, 13, 14, 15, 16], [[1, 1, 2], [1, 2, 2], [1, 2, 3], [1, 2, 4], [2, 2, 2], [2, 2, 3], [2, 2, 4], [2, 3, 3], [2, 3, 4], [2, 4, 4]])