An iterator returning all possible partitions of a list in Java

Question

Given a list of \$n\$ objects and an integer \$k \in \{ 1, 2, \dots, n \}\$, this iterator generates all possible ways of partitioning the elements in the list into exactly \$k\$ disjoint, non-empty blocks (partitions).

There is exactly \$S(n, k)\$ such partitions, see https://en.wikipedia.org/wiki/Stirling_numbers_of_the_second_kind.

My code is as follows:

PartitionIterable.java:

package net.coderodde.util;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;
import java.util.NoSuchElementException;

/**
 * This class implements an {@code Iterable} over all partitions of a given 
 * list.
 * 
 * @author Rodion "rodde" Efremov
 * @version 1.6 (Feb 14, 2016 a.k.a. Friend Edition)
 * @param <T> The actual element type.
 */
public class PartitionIterable<T> implements Iterable<List<List<T>>> {

    private final List<T> allElements = new ArrayList<>();
    private final int blocks;

    public PartitionIterable(List<T> allElements, int blocks) {
        checkNumberOfBlocks(blocks, allElements.size());
        this.allElements.addAll(allElements);
        this.blocks = blocks;
    }

    @Override
    public Iterator<List<List<T>>> iterator() {
        return new PartitionIterator<>(allElements, blocks);
    }

    private void checkNumberOfBlocks(int blocks, int numberOfElements) {
        if (blocks < 1) {
            throw new IllegalArgumentException(
                    "The number of blocks should be at least 1, received: " +
                    blocks);
        }

        if (blocks > numberOfElements) {
            throw new IllegalArgumentException(
                    "The number of blocks should be at most " +
                    numberOfElements + ", received: " + blocks);
        }   
    }

    private static final class PartitionIterator<T> 
    implements Iterator<List<List<T>>> {

        private List<List<T>> nextPartition;
        private final List<T> allElements = new ArrayList<>();
        private final int blocks;

        private final int[] s;
        private final int[] m;
        private final int n;

        PartitionIterator(List<T> allElements, int blocks) {
            this.allElements.addAll(allElements);
            this.blocks = blocks;
            this.n = allElements.size();

            s = new int[n];
            m = new int[n];

            if (n != 0) {
                for (int i = 0; i < n - blocks + 1; ++i) {
                    s[i] = 0;
                    m[i] = 0;
                }

                for (int i = n - blocks + 1; i < n; ++i) {
                    s[i] = m[i] = i - n + blocks;
                }

                loadPartition();
            }
        }

        @Override
        public boolean hasNext() {
            return nextPartition != null;
        }

        @Override
        public List<List<T>> next() {
            if (nextPartition == null) {
                throw new NoSuchElementException("No more partitions left.");
            }

            List<List<T>> partition = nextPartition;
            generateNextPartition();
            return partition;
        }

        private void loadPartition() {
            nextPartition = new ArrayList<>(blocks);

            for (int i = 0; i < blocks; ++i) {
                nextPartition.add(new ArrayList<>());
            }

            for (int i = 0; i < n; ++i) {
                nextPartition.get(s[i]).add(allElements.get(i));
            }
        }

        private void generateNextPartition() {
            for (int i = n - 1; i > 0; --i) {
                if (s[i] < blocks - 1 && s[i] <= m[i - 1]) {
                    s[i]++;
                    m[i] = Math.max(m[i], s[i]);

                    for (int j = i + 1; j < n - blocks + m[i] + 1; ++j) {
                        s[j] = 0;
                        m[j] = m[i];
                    }

                    for (int j = n - blocks + m[i] + 1; j < n; ++j) {
                        s[j] = m[j] = blocks - n + j;
                    }

                    loadPartition();
                    return;
                }
            }

            nextPartition = null;
        }
    }

    public static void main(String[] args) {
        List<String> list = Arrays.asList("A", "B", "C", "D");
        int row = 1;

        for (int blocks = 1; blocks <= list.size(); ++blocks) {
            for (List<List<String>> partition : 
                    new PartitionIterable<>(list, blocks)) {
                System.out.printf("%2d: %s\n", row++, partition);
            }
        }
    }
}

For example, all partitions of the set \$\{ A, B, C, D \}\$ are

 1: [[A, B, C, D]]
 2: [[A, B, C], [D]]
 3: [[A, B, D], [C]]
 4: [[A, B], [C, D]]
 5: [[A, C, D], [B]]
 6: [[A, C], [B, D]]
 7: [[A, D], [B, C]]
 8: [[A], [B, C, D]]
 9: [[A, B], [C], [D]]
10: [[A, C], [B], [D]]
11: [[A], [B, C], [D]]
12: [[A, D], [B], [C]]
13: [[A], [B, D], [C]]
14: [[A], [B], [C, D]]
15: [[A], [B], [C], [D]]

Is there anything to improve here?

Answer 1

You code looks really good but there are always ways for improvement :)

1. Add to main method printout of current \$k\$. For me output was unclear without it and I couldn't compare the output with "golden" table.
2. I would rename allElements to elements cause all- doesn't give any additional information.
3. Strings in IllegalArgumentException constructors may be changed to String.format("...text...", params...) it'll simplify i18n of such strings in the future
4. Use java.util.Arrays#fill for arrays filling instead of circles or even remove zeroing at all as Java Spec guarantees arrays zeroing after initialization.
5. Extract this initialization into a separate method (remember of SRP)
6. Make parameters' checks in checkNumberOfBlocks and PartitionIterator consistent (or even remove it from PartitionIterator at all as it's private impl)
7. Split generateNextPartition into some separate methods with SRP and appropriate names. E.g. for me is unclear how generateNextPartition works now. Remember empty lines in a method code is the sign of breaking SRP and is like "code smell" in the most cases

Answer 2

Looks good. For debugging and understanding I'd also print the block size in the test loop, that makes checking the results a bit more straightforward.

The initialisation of allElements / lists in general could be fully done in the constructor instead of splitting it up, i.e.:

    this.allElements = new ArrayList<>(allElements);

Copying over all the elements seems a bit wasteful, especially since that list is never modified. I'd argue that since you're only allowing Lists to be passed in the caller should have to provide an efficient list instead of you copying a lot internally. Same goes for modification: You're not going to modify the incoming list, so assume that the list is not changing underneath you and let the caller copy if necessary.

The same argument could be made for the generated partitions, but then again, the hassle of reusing elements is possibly not worth it for clarity. That said, the process of generateNextPartition and loadPartition looks a bit weird since it could very well be done in one step?

Removing the n != 0 condition in the PartitionIterator constructor doesn't change the behaviour, because the blocks argument is already checked outside of the iterator class.

Also initialising to zero doesn't do much there as both arrays are allocated right before that and will default to zero anyway.

That also reminds me that you should post the test cases as well. You do have some, right?