2015-02-12

How to generate all subsets of size k in Java

This blog post shows a Java class to generate all subsets of size k of the set {0, 1, ..., n - 1}, in lexicographic order. The code uses O(k) memory, and it doesn't store multiple subsets at the same time in memory. This is achieved by implementing the Iterable interface.

import java.util.Iterator;
import java.util.NoSuchElementException;

/**
 * Generate all subsets of {0, ..., n-1} of size k, in lexicographically
 * increasing order.
 */
public class Fixsub implements Iterator<int[]>, Iterable<int[]> {
  private int n;
  private int k;
  private int[] a;
  public Fixsub(int n, int k) {
    assert(k >= 0);
    assert(n >= k);
    this.n = n;
    this.k = k;
  }
  private int findIdxToIncrease() {
    int i;
    for (i = k - 1; i >= 0 && a[i] == n - k + i; --i) {}
    return i;
  }
  @Override public Iterator<int[]> iterator() {
    return this;
  }
  /**
   * Always returns the same array reference, the caller is responsible for
   * making a copy. The caller shouldn't modify the array elements.
   */
  @Override public int[] next() {
    if (a == null) {
      a = new int[k];
      for (int i = 0; i < k; ++i) {
        a[i] = i;
      }
    } else {
      int i = findIdxToIncrease();
      if (i < 0) throw new NoSuchElementException();
      for (++a[i++]; i < k; ++i) {
        a[i] = a[i - 1] + 1;
      }
    }
    return a;
  }
  @Override public boolean hasNext() {
    return a == null || findIdxToIncrease() >= 0;
  }
  @Override public void remove() {
    throw new UnsupportedOperationException();
  }

  public static void main(String[] args) {
    for (int[] p : new Fixsub(7, 4)) {
      StringBuilder sb = new StringBuilder();
      for (int i = 0; i < p.length; ++i) {
        if (i != 0) sb.append(',');
        sb.append(p[i]);
      }
      sb.append('\n');
      System.out.print(sb);
    }
  }
}

No comments: