python - 满足特定条件的列表列表的所有组合

Question

我希望获得列表列表的所有可能组合，但有一些特定条件。必须满足的条件:

1. 我假设输入列表中的所有子列表都具有相同的长度 N
2. 我只能为每列选择一个特定的金额
3. 输出应仅包含大小为 N 的列表
4. 每行我只能选择一个值(列)

例子:

# Columns 0 1 2
lst  =  [[1,2,3],
         [4,5,6],
         [7,8,9]]

choose = [2,1,0] # here choose only two from first column and one from middle column

为了帮助跟踪列，我修改了输入列表并将每个项目封装在一个对象(值，column_index)中。

>> lst = encapsulate(lst)
[[(1, 0), (2, 1), (3, 2)],
 [(4, 0), (5, 1), (6, 2)],
 [(7, 0), (8, 1), (9, 2)]]
>> combine(lst, choose)
[[(1, 0), (4, 0), (8, 1)],
 [(1, 0), (5, 1), (7, 0)],
 [(2, 1), (4, 0), (7, 0)]]

我有一个适用于较小列表的解决方案，但我的目标是在 20x12 矩阵上运行它。有什么可行的方法可以解决这个问题吗？

我目前的解决方案:

class Encapsulation:
    def __init__(self, value, column_index):
        self.value = value
        self.column_index = column_index

    def __repr__(self):
        return "(%d, %d)" % (self.value, self.column_index)


def combine(L, choose):
    combinations = []

    choose_n = len(choose)

    def _combine(terms, accum):
        last = (len(terms) == 1)
        n = len(terms[0])
        for i in range(n):
            item = accum + [terms[0][i]]
            if last:
                if can_choose(item, choose, choose_n):
                    combinations.append(item)
            else:
                if can_choose(item, choose, choose_n):
                    _combine(terms[1:], item)

    _combine(L, [])
    return combinations

def encapsulate(lst):
    outlist = []
    for j in range(0, len(lst)):
        new_l = []
        for i in range(0, len(lst[j])):
            new_l.append(Encapsulation(lst[j][i], i))
        outlist.append(new_l)
    return outlist

def can_choose(l, _choose, n):
    counts = [0]*n

    for _item in l:
        counts[_item.column_index] += 1
        if counts[_item.column_index] > _choose[_item.column_index]:
            return False

    return True


if __name__ == "__main__":
    lst = [[1,2,3],
           [4,5,6],
           [7,8,9]]

    choose = [2,1,0]

    lst = encapsulate(lst)
    assert(sum(choose) == len(lst) and len(choose) == len(lst[0]))
    combinations = combine(lst, choose)
    print(combinations)

编辑 1:

我设法稍微加快了实现速度。但对于大型矩阵仍然不可行。

def product(*args, **kwds):
    "Alternative fast implementation of product for python < 2.6"
    choose = kwds.get('choose', [])
    choose_n = len(choose)

    def cycle(values, uplevel):
        for prefix in uplevel:       # cycle through all upper levels
            counts = [0]*choose_n

            for _item in prefix:
                counts[_item.column_index] += 1

            for current in values:   # restart iteration of current level
                i = current.column_index
                if counts[i]+1 <= choose[i]:
                    yield prefix + (current,)

    stack = iter(((),))
    for level in tuple(map(tuple, args)):
        stack = cycle(level, stack)  # build stack of iterators
    return stack

像这样使用

lst  =  [[1,2,3],
         [4,5,6],
         [7,8,9]]

choose = [2,1,0]
product(*lst, choose=choose)

Answer 1

这是一个使用 itertools 的解决方案，它将生成可能的输出，而无需每个解决方案的可能排列。如果您还需要这些排列，只需生成每个输出的所有排列即可。

我首先构建列，然后生成每列所请求的值的所有可能组合的生成器。最后，我们制作这些组合的产品:

from itertools import combinations, product, chain

lst  =  [[1,2,3],
         [4,5,6],
         [7,8,9]]

choose = [2,1,0]


columns = list(zip(*lst))
# [(1, 4, 7), (2, 5, 8), (3, 6, 9)]

if sum(choose) != len(columns[0]):
    raise ValueError(f'You must choose {len(columns[0])} values')

combinations = [combinations(columns[index], r=number) for index, number in enumerate(choose)]

for comb in product(*combinations):
    out = list(chain.from_iterable(comb))
    print(out)

输出:

[1, 4, 2]
[1, 4, 5]
[1, 4, 8]
[1, 7, 2]
[1, 7, 5]
[1, 7, 8]
[4, 7, 2]
[4, 7, 5]
[4, 7, 8]