python: генерация вариантов слагаемых суммы с помощью itertools (или иных библиотек)

Question

Подскажите как быстро решить следующую задачу:

есть сумма S и n слагаемых, причём первое и последнее слагаемые могут быть равными 0, а все остальные слагаемые только больше 0 (от 1 и выше)

необходимо сформировать все возможные варианты

например для S = 2 и n = 3 будут такие варианты:

0 1 1
0 2 0
1 1 0

создал такой код:

# вычислить варианты
def calculate_variants(variants, variant, s, n_max, n):
    if n == 1:
        variants.append(variant + [s])
    else:
        for i in range(0 if n == n_max else 1, s + 1):
            calculate_variants(variants, variant + [i], s - i, n_max, n - 1)

# получить варианты
def get_variants(size, data):
    variants = []
    calculate_variants(variants, [], s, n, n)
    return variants

код конечно немного кривоватый, но рабочий

но хотелось бы более производительный алгоритм

подскажите можно ли это как то с помощью библиотек типа itertools сделать, чтоб и код был покрасивше :) и попроизводительнее

Answer 1

Основное время занимает организация хранения - ведь вариантов очень много, их число растёт экспоненциально. Если вместо принта активировать елд, то время работы без сохранения (а варианты при этом генерируются) и с сохранением результатов в двумерный список (v = [x for x in compo(6, 3)]) отличается раз в 30.

Поэтому, если есть возможность - обработать вариант и забыть про него. Если всё надо сохранять в файл - возможно, запись каждого варианта по очереди даст какой-то выигрыш.

Код этот адаптирован из моего ответа на EnSO, там же приводится сравнение с компилируемыми языками.

Генерируется комбинация из n чисел, в сумме дающих s-(n-2), затем к средним элементам добавляются единицы

def compo(n, s):
    x = [0] * n
    s = s - n + 2
    x[0] = s
    adder = [1] * n
    adder[0] = 0
    adder[-1] = 0

    while True:
        print([a+b for a,b in zip(x,adder)])
        #yield([a+b for a,b in zip(x,adder)])
        v = x[-1]
        if (s==v ):
            break
        x[-1] = 0
        j = -2
        while (0==x[j]):
            j -= 1
        x[j] -= 1
        x[j+1] = 1 + v

Answer 2

Наверное оно

from itertools import product

def compo(n, s):
    return [x for x in product(range(s+1), repeat=n) if sum(x) == s and 0 not in x[1:-1]]

На выходе будет список кортежей. С ними уже делайте, что нужно.

Answer 3

Не быстрый вариант, но как источник возможных идей:

def variants(s, n):
    for c in itertools.combinations(range(s + 1), n - 1):
        yield np.diff((0, ) + c + (s, ))


for c in variants(4, 3):
    print(*c)

Answer 4

Рекурсивная генерация с numpy. numpy экономит память и очень быстро работает при копировании частей таблицы (нужно для заполнения и кеширования). Таблица с комбинациями выделяется один раз в начале работы. В кеше хранятся указатели на куски таблицы. Сам массив cache требует очень мало памяти. order='F' улучшает скорость копирования кешированных кусков в три раза.

def variants(s, n):
    result = np.zeros((math.comb(s + 1, n - 1), n), dtype=np.uint8, order='F')
    cache = [[None] * n for _ in range(s + 1)]

    def v_range(s, i):
        if i == 0:
            return 0, s - (n - 2) + 1
        if i == n - 1:
            return s, s + 1
        return 1, s - (n - i - 2) + 1

    def fill(p, s, i):
        if i == n:
            return p + 1

        if cache[s][i] is None:
            q = p
            for v in range(*v_range(s, i)):
                qq = fill(q, s - v, i + 1)
                result[q:qq, i] = v
                q = qq
            cache[s][i] = p, q
            return q

        cp, cq = cache[s][i]
        q = p + (cq - cp)
        result[p:q, i:n] = result[cp:cq, i:n]
        return q

    q = fill(0, s, 0)
    assert q == result.shape[0]
    return result

 s n        число      время
       комбинаций     работы

26 7       296010     0.0035
50 7     18009460     0.0672
97 7   1052618392     3.5804

Последний результат требует около 7Gb. Отдельного кеша нет - оптимальный результат.

P.S. Это легко переносится на C++. Интересно, кто будет быстрее: Python или C++?

БЛОГ НА HUSL

python: генерация вариантов слагаемых суммы с помощью itertools (или иных библиотек)

Ответы (4 шт):