組み合わせを列挙する

SRMとかやってると組み合わせをみる必要があったりして、そのたびに同じようなコードを書いていたので、なんとか出来ないかと思って書いてみた。

import java.util.*;
public class Combinations<T> implements Iterator{

    private List<List<T>> combinations;
    private List<T> list;
    private int[] index;
    private boolean[] visited;
    private int r;
    private Iterator<List<T>> iterator;

    public Combinations(T[] array, int r){
        this.list = Arrays.asList(array);
        this.index = new int[r];
        this.visited = new boolean[array.length];
        this.combinations = new ArrayList<List<T>>();
        this.r = r;
        this.compute(0);
        this.iterator = this.combinations.iterator();
    }

    private void compute(int n){
        if(n == r){
            List<T> combination = new ArrayList<T>();
            for(int i = 0;i < this.index.length;i++){
                combination.add(list.get(index[i]));
            }
            combinations.add(combination);
        }
        else{
            for(int i = 0;i < this.list.size();i++){
                if(n == 0 || !this.visited[i] && index[n - 1] < i){
                    this.visited[i] = true;
                    this.index[n] = i;
                    this.compute(n + 1);
                    this.visited[i] = false;
                }
            }
        }
    }

    public List<List<T>> getCombinations(){
        return this.combinations;
    }

    public int count(){
        return this.combinations.size();
    }

    public List<T> next(){
        return this.iterator.next();
    }

    public boolean hasNext(){
        return this.iterator.hasNext();
    }

    public Iterator<List<T>> iterator(){
        return this.iterator;
    }

    public void remove(){
    }

    //test
    public static void main(String[] argv){
        String[] test = new String[]{
            "a", "b", "c", "d", "e", "f", "g"
        };
        Combinations<String> c = new Combinations<String>(test, 2);
        Iterator<List<String>> i = c.iterator();
        while(i.hasNext()){
            List<String> t = i.next();
            System.out.println(t.get(0) + " " + t.get(1));
        }
        System.out.println(c.count() + " petterns.");
    }

}

結果:

$ java Combinations
a b
a c
a d
a e
a f
a g
b c
b d
b e
b f
b g
c d
c e
c f
c g
d e
d f
d g
e f
e g
f g
21 petterns.

多分合ってる。
ちょっと大きい数字でやるとすぐ計算時間が膨大になる気がする。

追記：
100C100とかやってもすぐ帰ってきてほしいので、改良した。

import java.util.*;
public class Combinations<T> implements Iterator{

    private List<List<T>> combinations;
    private List<T> list;
    private int[] index;
    private boolean[] visited;
    private int r;
    private boolean overHalf;
    private Iterator<List<T>> iterator;

    public Combinations(T[] array, int r) throws IllegalArgumentException {
        if(array.length < 1 || r < 1 || array.length < r){
            throw new IllegalArgumentException();
        }

        this.combinations = new ArrayList<List<T>>();
        this.list = Arrays.asList(array);
        this.r = r;
        if(this.r == array.length){
            this.combinations.add(list);
        }
        else{
            if(this.r > list.size() / 2){
                this.r = list.size() - this.r;
                this.overHalf = true;
            }

            this.index = new int[this.r];
            this.visited = new boolean[list.size()];
            this.compute(0);
        }

        this.iterator = this.combinations.iterator();
    }

    private void compute(int n){
        if(n == this.r){
            List<T> combination = new ArrayList<T>();
            if(overHalf){
                for(int i = 0;i < this.list.size();i++){
                    boolean skip = false;
                    for(int j = 0;j < this.index.length;j++){
                        if(i == this.index[j]){
                            skip = true;
                        }
                    }
                    if(skip){
                        continue;
                    }
                    combination.add(list.get(i));
                }
            }
            else{
                for(int i = 0;i < this.index.length;i++){
                    combination.add(list.get(index[i]));
                }
            }
            this.combinations.add(combination);
        }
        else{
            for(int i = 0;i < this.list.size();i++){
                if(n == 0 || !this.visited[i] && index[n - 1] < i){
                    this.visited[i] = true;
                    this.index[n] = i;
                    this.compute(n + 1);
                    this.visited[i] = false;
                }
            }
        }
    }

    public List<List<T>> getCombinations(){
        return this.combinations;
    }

    public static int count(int n, int r) throws IllegalArgumentException{
        if(n < 1 || r < 1 || n < r){
            throw new IllegalArgumentException();
        }

        int res = 1;
        for(int i = n;i > r;i--){
            res *= i;
        }
        for(int i = n - r;i > 1;i--){
            res /= i;
        }
        return res;
    }

    public List<T> next(){
        return this.iterator.next();
    }

    public boolean hasNext(){
        return this.iterator.hasNext();
    }

    public Iterator<List<T>> iterator(){
        return this.iterator;
    }

    public void remove(){
    }

    //test
    public static void main(String[] argv){
        String[] test = new String[]{
            "a", "b", "c", "d", "e", "f", "g"
        };
        int r = 5;
        Combinations<String> c = new Combinations<String>(test, r);
        Iterator<List<String>> it = c.iterator();
        while(it.hasNext()){
            List<String> comb = it.next();
            for(int i = 0;i < comb.size();i++){
                System.out.print(comb.get(i) + " ");
            }
            System.out.println("");
        }
        System.out.println(Combinations.count(test.length, r) + " petterns.");
    }
}

100C99っていうのは100C1と計算結果的には同じなので、rが半分より多い場合は「選ばないものを選ぶ」という感じになっている。100C99とかやっても割と早く帰ってくると思う。