Generate integer partitions and compositions

partitions.py

#!/usr/bin/env python3

def main():
    print("partitions of 5:")
    for p in partitions(5):
        print(p)
    print()

    print("compositions of 5:")
    for c in compositions(5):
        print(c)


def partitions(n):
    # This generates the partitions of n in reverse lexicographical order
    # (higher numbers come first in each partition).
    if n == 0:
        # Trivially, there is only the empty partition of 0.
        yield ()
    else:
        # Start by generating the partitions of n-1.
        for p in partitions(n - 1):
            # We can always create a unique partition by appending 1 to the end
            # of partitions of n-1. This will always be canonical becacuse 1
            # is the lowest possible number of any part and we're putting it
            # last.
            yield p + (1,)
            if len(p) == 1 or (len(p) > 1 and p[-2] > p[-1]):
                # Similarly, we can create a new partition by adding 1 to any
                # constituent element of each partitions of n-1. However, we
                # need to avoid duplicates. To do so, we ensurer that the output
                # is canonical. This is achieved by only adding the the last
                # element and only if the penultimate element (if it exists) is
                # strictly greater than the last element. This keeps the output
                # in reverse lexicographical order.
                yield p[:-1] + (p[-1] + 1,)

def compositions(n):
    # NOTE: It's natural to think that we might be able to bootstrap from the
    # partitions generator above. For example, by inserting a 1 into each
    # position of compositions of n - 1 (avoiding multiple insertions into runs)
    # and then adding a 1 to each composition of n - 1. Unfortunately, this
    # results in duplicate compositions due to some addition operations
    # resulting in the same output as other addition operations. (You can easily
    # see that these will never be equivalent to the 1 insertions since those
    # always change the size of the partition, so there will be no overlap.)
    #
    # Instead, we generate k-compositions of n (n into exactly k parts) for
    # k = 1 to n.
    for k in range(1, n+1):
        yield from k_compositions(n, k)

def k_compositions(n, k):
    # Generate all k-compositions in increasing lexicographical order. See
    # https://stackoverflow.com/a/69696615 for notes.
    if k > n:
        # Impossible (without weak compositions)
        return
    if k == 0:
        # The loop below assumes we have at least a zeroth element. Special
        # casing here lets us short-circuit that.
        yield ()
        return
    c = (1,) * (k - 1) + (n - k + 1,)
    yield c
    while c[0] < n - k + 1:
        # We're looking for the rightmost index where all values to the right of
        # it are 1 and whose value is greater than 1. Call this value m, i.e.,
        # c[idx] = m, and m > 1 and all values to the right of idx are 1. In
        # other words, everything to the right of and including idx makes up the
        # highest lexicographic composition of (m + n - k + 1) into (n-k) parts.
        # This means that we've exhausted all compositions of (m + n - k + 1)
        # and must reduce it to compositions of (m + n - k) into (n - k) parts.
        # We know this is possible because m > 1, i.e., we won't end up with a
        # weak composition by doing this. To keep the entire sequence a
        # k-composition of n, we need to add one to the value at (idx - 1). At
        # the same time, we want to start with the _lowest_ lexicographical
        # composition of (m + n - k) into (n - k) parts, so we need to keep the
        # rest of the right subsequence as 1s and put ((m + n - k) - (n - k - 1))
        # = (m - 1) into the rightmost position. We repeat this process until we
        # reach the highest lexicographical value (highest possible value in the
        # first index of the k-composition).
        idx = k - 1
        while c[idx] == 1:
            idx -= 1
        assert idx > 0 # By construction, we know we are not at the last composition
        m = c[idx]
        # Annoyingly, we have to perform the update operation in multiple steps.
        # This lets us avoid special casing where idx == k - 1.
        #   original     incremented      tail incl idx
        #   |            |                |
        #   v            v                v
        c = c[:idx-1] + (c[idx-1] + 1,) + c[idx+1:]
        #   original   idx   tail right of idx
        #   |          |     |
        #   v          v     v
        c = c[:idx] + (1,) + c[idx+1:]
        #   up to last last
        #   |          |
        #   v          v
        c = c[:k-1] + (m-1,)
        yield c

if __name__ == "__main__":
    main()

partitions.rs

fn main() {
    println!("partitions of 5:");
    partitions(5, &mut |p| {
        println!("{p:?}");
    });
    println!("compositions of 5:");
    compositions(5, &mut |p| {
        println!("{p:?}");
    });
}

fn partitions<F>(n: u8, f: &mut F)
where
    F: FnMut(&[u8]),
{
    fn rec(n: u8, f: &mut dyn FnMut(&mut Vec<u8>), p: &mut Vec<u8>) {
        if n == 0 {
            f(p);
        } else {
            let mut f_rec = |p: &mut Vec<u8>| {
                p.push(1);
                f(p);
                p.pop();
                if p.len() == 1 || (p.len() > 1 && p[p.len() - 2] > p[p.len() - 1]) {
                    let idx = p.len() - 1;
                    p[idx] += 1;
                    f(p);
                    p[idx] -= 1;
                }
            };
            rec(n - 1, &mut f_rec, p);
        }
    }
    rec(n, &mut |p| f(p), &mut Vec::with_capacity(n as usize));
}

fn compositions<F>(n: u8, f: &mut F)
where
    F: FnMut(&[u8]),
{
    for k in (1..=n).rev() {
        k_compositions(n, k, f);
    }
}

fn k_compositions<F>(n: u8, k: u8, f: &mut F)
where
    F: FnMut(&[u8]),
{
    if k > n {
        return;
    }
    if k == 0 {
        f(&[]);
        return;
    }
    let mut c = vec![1; k as usize];
    let last_idx = c.len() - 1;
    c[last_idx] = n - k + 1;
    f(&c);
    while c[0] < n - k + 1 {
        let mut idx = (k - 1) as usize;
        while c[idx] == 1 {
            idx -= 1;
        }
        assert!(idx > 0);
        let m = c[idx];
        c[idx - 1] += 1;
        c[idx] = 1;
        c[(k - 1) as usize] = m - 1;
        f(&c);
    }
}