Enumerates integer partitions in Gray-code order

gray_coded_partitions.py

# An implementation of Dr. Carla Savage's algorithm described in her
# 1989 paper "Gray code sequences of partitions"
# https://www.sciencedirect.com/science/article/abs/pii/0196677489900072

# Referred to corrections made in a lecture by Dr. Sriram Pemmaraju
# https://homepage.cs.uiowa.edu/~sriram/196/fall01/lecture7.pdf

# Implemented by Tom Quinn (https://thquinn.github.io/)
# I've made my own corrections, see the comments.

# Returns a Gray-code-order (minimal-change) list of partitions of `n`
# into parts of size at most `k`. Starts at the lexicographical min of the
# partitions. Ends at the lex. max unless (n, k) == (6, 4).
def L(n, k):
    if n == 0:
        return [[]]
    if k == 0:
        return []
    if k == 1:
        return [[1] * n]
    if (n, k) in L_special_cases:
        case = [expand_special_case_string(s) for s in L_special_cases[(n, k)]]
        return [i for l in case for i in l]
    if n < k:
        return L(n, n)
    if n >= 2 * k:
        return (
            M(n, k - 1) +
            prepend_all([k, k - 1], L(n - 2 * k + 1, k - 1))[::-1] +
            prepend_all([k, k], L(n - 2 * k, k))
        )
    if n == 2 * k - 1:
        return (
            M(n, k - 1) +
            [[k, k - 1]]
        )
    if n == 2 * k - 2:
        # correction to the paper from Dr. Pemmaraju here:
        # second term is reversed, third term is NOT reversed
        return (
            L(n, k - 2) +
            prepend_all([k - 1], L(n - k + 1, k - 3))[::-1] +
            prepend_all([k], L(n - k, k - 3)) + # (2, 2) -> (0, -1), added L(2, 2) special case
            [[k - 1, k - 2, 1]] +
            [[k - 1, k - 1]] +
            [[k, k - 2]]
        )
    return (
        L(n, k - 2) +
        prepend_all([k - 1], L(n - k + 1, n - k + 1))[::-1] +
        prepend_all([k], L(n - k, n - k))
    )

def M(n, k):
    if k == 0:
        return []
    if k == 1:
        return [[1] * n]
    if (n, k) in M_special_cases:
        case = [expand_special_case_string(s) for s in M_special_cases[(n, k)]]
        return [i for l in case for i in l]
    if 2 * k + 1 <= n and n < 3 * k - 1:
        return (
            L(n, k - 1) + 
            prepend_all([k], L(n - k, k - 1))[::-1] +
            prepend_all([k + 1], L(n - k - 1, k - 2)) +
            prepend_all([k + 1, k - 1], L(n - 2 * k, k - 2))[::-1] +
            prepend_all([k, k], L(n - 2 * k, k))
        )
    return (
        L(n, k - 1) +
        prepend_all([k], M(n - k, k - 1))[::-1] +
        prepend_all([k + 1], L(n - k - 1, k - 1)) +
        prepend_all([k, k, k - 1], L(n - 3 * k + 1, k - 1))[::-1] +
        ([] if n == 3 * k - 1 else prepend_all([k, k, k], L(n - 3 * k, k)))
    )
        
def prepend_all(elements, partitions):
    return [elements + partition for partition in partitions]

def expand_special_case_string(s):
    if s[-1] != ')':
        return [[int(c) for c in s]]
    fn_index = [c.isalpha() for c in s].index(True)
    prefix_elements = [int(c) for c in s[:fn_index]]
    fn = s[fn_index:]
    arg_n = int(fn[fn.find('(')+1:fn.find(',')])
    arg_k = int(fn[fn.find(',')+1:-1])
    if fn.startswith('L'):
        return prepend_all(prefix_elements, L(arg_n, arg_k))
    if fn.startswith('RL'):
        return prepend_all(prefix_elements, L(arg_n, arg_k))[::-1]
    if fn.startswith('M'):
        return prepend_all(prefix_elements, M(arg_n, arg_k))
    if fn.startswith('RM'):
        return prepend_all(prefix_elements, M(arg_n, arg_k))[::-1]
    
L_special_cases = {
    # this (2, 2) case appears necessary, otherwise it hits L3.d.
    (2, 2): ['11', '2'],
    (6, 4): ['111111', '21111', '3111', '2211', '222', '321', '33', '42', '411'],
    # excluded the (6, 6) special case, isn't needed
    (12, 7): [
        '111111111111', '21111111111', '3111111111', '2211111111', '222111111', '22221111',
        '2222211', '222222', '322221', '3222111', '32211111', '321111111', '33111111',
        '3321111', '333111', '332211', '33222', '33321',
        # this case was cut off here, my continuation
        '3333', '4332', '4422', '44211', '441111', '4311111', '432111', '43311', '43221',
        '42222', '422211', '4221111', '42111111', '411111111', '51111111', '5211111', '531111',
        '522111', '52221', '5322', '53211', '5331', '4431', '444', '543', '5421', '54111',
        '5511', '552', '642', '633', '732', '7221', '72111', '711111', '6111111', '621111',
        '63111', '62211', '6222', '6321', '6411', '7311', '741', '651', '66', '75'
    ],
    (14, 4): [
        '11111111111111', '2111111111111', '221111111111', '22211111111', '2222111111',
        '222221111', '22222211', '2222222', '3222221', '422222', '4222211', '32222111',
        '322211111', '42221111', '422111111', '3221111111', '4211111111', '41111111111',
        '311111111111', '32111111111', '3311111111', '431111111', '44111111', '4421111',
        '442211', '44222', '43322', '433211', '4331111', '43211111', '4322111', '432221',
        '332222', '3322211', '33221111', '332111111', '33311111', '3332111', '333221',
        '333311', '33332', '43331', '4433', '44321', '443111', '44411', '4442'
    ],
    (15, 5): [
        'M(15,4)', '5433', '5442', '54411', '543111', '54321', '54222', '542211',
        '5421111', '54111111', '55L(5,5)'
    ],
}
M_special_cases = {
    (11, 5): [
        'L(11,4)', '533', '542', '5411', '53111', '5321', '5222', '52211', '521111',
        '5111111', '6L(5,3)', '641', '551'
    ],
    (13, 6): [
        'L(13,5)', '6RL(7,5)', '7L(6,3)', '742', '7411', '751', '661'
    ],
    (18, 4): [
        'L(18,3)', '433332', '4333221', '4333311', '43332111', '433311111', '4332L(6,2)',
        '44RL(10,2)', '4311111111111', '432111111111', '43311111111', '4322L(7,2)',
        '4RL(14,2)', '5L(13,3)', '443331', '44433',
        # discontinuity in the case here, my continuation
        '444321', '444222', '4442211', '44421111', '444111111', '4431111111', '443211111',
        '44331111', '44322111', '4432221', '443322', '4433211', '4443111', '444411', '44442'
    ],
    (20, 5): [
        'L(20,4)', '5RM(15,4)', '6L(14,2)', '6322RL(7,2)', '63311111111', '632111111111',
        '6311111111111', '64L(10,2)', '6332RL(6,2)', '6333L(5,3)', '643RL(7,3)', '644L(6,3)',
        '64442', '644411', '554411', '55442', '554RL(6,3)', '555L(5,5)'
    ],
}