This is a fun exercise in combinatorics. I vaguely remember that this is pretty well-known and standard, but I donโt recall the name, so Iโll instead re-derive the construction here.
edit: @Syx_Pek provided the name down-thread, itโs called combinatorial number system.
Originally I described a re-derivation with code, but read wikipedia instead, itโs clearer than what I wrote.
click here to expand my pre-edit off-the-cuff derivation
First of all, suppose you had your array inv_co(k) = filter(y->count_ones(y)==k, 1: (-1 >>> 1))
, and youโre given a specific N
, such that count_ones(N)=k
. What is the index of N
in this list inv_co(k)
?
Well, it is one plus the number of smaller entries. This is easy to compute:
struct BitsIter
u::Int64
end
Base.iterate(b::BitsIter) = iterate(b, b.u)
Base.iterate(b::BitsIter, u) = (u == 0 ? nothing : (trailing_zeros(u), Base._blsr(u)))
function index(N)
res = 1
for (nbits, bitpos) in enumerate(BitsIter(N))
#there are nbits many bits we need to distribute over bitpos many places
res += binomial(bitpos, nbits)
end
res
end
We can reality-check that:
julia> inv_count_ones(n, z) = filter(y->count_ones(y)==n, 1:z)
inv_count_ones (generic function with 1 method)
julia> length(inv_count_ones(5, 1000))
252
julia> map(index, inv_count_ones(5, 1000)) == 1:length(inv_count_ones(5, 1000))
true
OK, now we need to invert the index
function (you want random access to your iterator / lazy collection of preimages, right?). That is: We are given idx and K, and want z
such that count_ones(z) == K
and index(z) == idx
.
I think the construction is best explained by giving a (slowish recursive) implementation:
function inv_index0(k, idx)
k == 1 && return 1<<idx
pos = findfirst(p -> binomial(p, k) > idx, 1:64)
rem = idx - binomial(pos-1, k)
@show k, pos, idx
return (1<<(pos-1)) | inv_index0(k-1, rem)
end
inv_index(k, idx) = inv_index0(k, idx-1)
We can reality check that:
julia> for i = 1:1000
z = rand(1:(1<<35))
k = count_ones(z)
if z != inv_index(k, index(z))
@show z, k, i
@assert false
end
end
Remaining parts of the exercise: Make a loop out of the tail-recursion, deal with integer overflows, make it fast (lookup table), optionally make a fast iterator, optionally make a fast container.