设置每个元素的 itertools 产品重复值

Set itertools product repeat value per element

如果我运行代码

import itertools

products = itertools.product([0,1],repeat=3)
print(list(products))

我得到以下输出:

[(0, 0, 0), (0, 0, 1), (0, 1, 0), (0, 1, 1), (1, 0, 0), (1, 0, 1), (1, 1, 0), (1, 1, 1)]

但是,我只想重复 0 一次和 1 两次。换句话说,我想要以下输出:

[(0, 1, 1), (1, 0, 1), (1, 1, 0)]

我怎样才能做到这一点?


当然可以:

import itertools

products = itertools.permutations([0,1,1],3)
print(list(set(products)))

但在我的例子中有大量元素,因此在迭代之前调用 set 会由于内存问题而终止代码。

这可能对您的示例过于专业化,但是

>>> from collections import namedtuple
>>> t = namedtuple('t', 'v1 v2 v3')
>>> [tuple(t(1,1,1)._replace(**{x: 0})) for x in t._fields]
[(0, 1, 1), (1, 0, 1), (1, 1, 0)]

如果你只有 0 和 1,这会起作用:

from itertools import combinations

def gen(n0, n1):
    n = n0 + n1
    for c in combinations(range(n), n1):
        out = [0]*n
        for i in c:
            out[i]=1
        yield out

list(gen(1,2))

构建out的方法可能不是最优的,但想法是存在的,如果时间效率有问题,我会留给你改进它。

进一步概括:

def gen(n0, n1, n2):
    n12 = n1 + n2
    n = n0 + n12
    for c12 in combinations(range(n), n12):
        out = [0]*n
        for i in c12:
            out[i] = 1
        for c2 in combinations(c12, n2):
            out_ = out.copy()
            for i in c2:
                out_[i] = 2
            yield out_

同样,out_ 的构造可能不是最理想的。使用相同的想法,您可以继续嵌套越来越多的不同元素。如果你有更多可能的元素,深度变得手工嵌套很麻烦,你可以递归这个过程,这也是一个有趣的练习:

def gen(ns, elems=None, C=None, out=None):
    
    if elems is None:
        elems = list(range(len(ns)))
    else:
        assert len(elems) == len(ns)
        
    if out is None:
        N = 1
        for n in ns:
            N *= n
        out = [elems[0]]*N
        C = range(N)
    
    if len(ns) == 1:
        yield out
    
    else:
        n = ns[-1]
        e = elems[-1]
        
        for c in combinations(C,n):
            out_ = out.copy()
            for i in c:
                out_[i] = e
            C_ = [i for i in C if i not in c]
            yield from gen(ns[:-1], elems[:-1], C_, out_)