Search code examples
algorithmrustcombinatorics

Parse every k-element subset of an n-element vector


My goal is to perform an operation on every k-element subset of n elements stored in a vec.

My current approach uses this bit-manipulation to generate integers with k set bits in lexicographic order. I.e., given one such integer, it modifies it to the next one in order.

pub fn set_next_bit_permutation(v: &mut u128) {
    let t: u128 = *v | (*v-1);
    *v = (t + 1) | (((!t & (!t).wrapping_neg()) - 1) >> (v.trailing_zeros() + 1))
}

Then, I use positions of the set bits as the indices, pulling out each index like this. Say I'm using perm: u128 to store my current permutation. This code would be inside a loop that runs until I check every permutation:

perm_copy = perm;
while perm_copy > 0 {
  cur_index = perm.trailing_zeros();
  // do something with cur_index
  perm_copy ^= 1 << cur_index;
}
set_next_bit_permutation(perm);

While this works and is pretty fast (important because we go through this many times), we're limited to size 128. Also, it feels like a hack.

Questions:

  • Does Rust have anything built-in to handle fixed-size subsets of Vecs? Or, I'm not married to Vecs if some other container handles this.
  • If not, is there a better, more idiomatic way to accomplish this?
  • Or, how could I modify this to handle more than length 128? Arbitrarily large?

Solution

  • It's actually easy and efficient to do this with a sorted array of k indices.

    Lets sort in ascending order and produce arrays of indices in lexicographic order considering the highest index first.

    So we initialize the array with the lowest indexes: A=[0, 1, ... k-1]

    Then, to get the next array of subset indexes in order, just:

    1. Find the smallest index i such that i == k-1 or A[i+i] > A[i]+1
    2. Increment A[i]
    3. Set the preceding elements to [0, 1, ... i-1]

    When you increment A[k-1] to n, you're done.