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



如果我运行代码

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将由于内存问题而杀死代码。

如果你只有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_)

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

>>> 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)]

最新更新