如果我运行代码
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)]