条件を満たすリストの作成の効率化について

Question

前提・実現したいこと

python3において
以下の条件を満たすようなリストを作成したいです．
・リスト内の要素の総和がn
・ある要素の和が与えられた条件を満たしている
・条件は要素の和が与えられる．与えられないものは基準値以下
作成するリストの長さが大きく，リストの要素の総和が大きいとき場合時間がかかってしまう（当たり前ですが...）
効率良く条件に合うリストを作り出すことはできないでしょうか？
現在は，リストの候補（総和n，リストの長さl_lenのリスト）をすべて計算し，その中から条件に合うものを結果として出力しています．

該当のソースコード

from itertools import combinations_with_replacement
n=20  #リストの要素の総和
l_len = 5 #作成するリストの長さ
lvalue = 2 #基準値
check_dict = {frozenset([1,2,3]) : 12 , frozenset([0,2,4]) : 12 , frozenset([0,1]) : -1} #key:リストの要素のindex ， value:keyが示すリストの要素の和(-1の場合基準値以下ならOK)
result = [] #条件に合ったリストを格納するリスト
#リストの要素の総和がn，作成するリストの長さがl_lenのリストの総組み合わせ
for x in combinations_with_replacement(range(n+1), l_len-1):
    t = [ed - st for st, ed in zip((0,) + x, x + (n,))]
    #リストのチェック
    for k,v in check_dict.items():
        s = 0
        for index in k:
            s += t[index]
        if v == -1 and s > lvalue:
            break
        if v != -1 and s != v:
            break
    else:
        result.append(t)
result

今回の出力

[[0, 0, 4, 8, 8],
[0, 1, 4, 7, 8],
[0, 2, 4, 6, 8],
[1, 0, 4, 8, 7],
[1, 1, 4, 7, 7],
[2, 0, 4, 8, 6]]

Comments

[kairi003]

質問なのですが，[0,0,0,0,20]みたいな要素も条件を満たしますよね？

[kairi003]

すみません，条件の部分の理解が足りませんでした．大丈夫です．
ただcheck_dictを辞書型である意味がないので素直に二次元タプルあたりにしたほうがいいと思います．

Answer 1

こうなりました。

import time

n=20  #リストの要素の総和
l_len = 5 #作成するリストの長さ
lvalue = 2 #基準値
check_dict = {frozenset([1,2,3]) : 12 , frozenset([0,2,4]) : 12 , frozenset([0,1]) : -1} #key:リストの要素のindex ， value:keyが示すリストの要素の和(-1の場合基準値以下ならOK)
result = [] #条件に合ったリストを格納するリスト

def f(n, l_len, lvalue, check_dict, result, temp = None, level = 0):
    if temp is None:
        # 初回呼び出し時の初期化
        temp = [ 0 for _ in range(l_len)]
        remain = n
    else:
        remain = n - sum(temp)


    # 指定された添え字における候補を決定
    possible = [ True for _ in range(remain + 1) ]

    for k, v in check_dict.items():
        if level in k:
            c_sum = sum([ temp[_] for _ in k])
            if v == -1:
                # lvalueを超える組み合わせを要素
                for i in range(lvalue - c_sum+1, remain + 1):
                    possible[i] = False
            elif level == list(k)[-1]:
                # 合計値指定のfrozensetの最後の場合は候補は一つ(v - c_sum)。それ以外を除外
                for i in range(0, min(v - c_sum, remain + 1)):
                    possible[i] = False
                for i in range(v - c_sum+1, remain + 1):
                    possible[i] = False
            else:
                # 合計値指定のfrozensetの途中の場合は上限を超える値を除外
                for i in range(v - c_sum +1, remain + 1):
                    possible[i] = False

    if level == (l_len - 1):
        # リストの最後の場合は残りのみで判定
        if possible[remain]:
            temp[level] = remain
            result.append(temp.copy())
            temp[level] = 0
    else:
        # 候補を使いながら再起呼び出し
        for i in range(remain + 1):
            if possible[i]:
                temp[level] = i
                f(n, l_len, lvalue, check_dict, result, temp, level+1)
                temp[level] = 0

start = time.time()
f(n, l_len, lvalue, check_dict, result)
elapsed_time = time.time() - start
print(elapsed_time)
print(result)

私の環境では質問欄コードは約20msに対して、回答コードは約1msになります。