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

Posted 2023-02-06

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

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)

答案

[这里是使用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]