创建维N的所有向量集的优雅方法是什么,每个元素都是0到K之间的整数([0,K])。
我目前的代码是:
def nodes_init(n, k):
nodes = {}
e = np.identity(n)
nodes[tuple(np.zeros(n))] = 0
s = Set()
s.add(tuple(np.zeros(n)))
s_used = Set()
while len(s) != 0:
node = s.pop()
if node in s_used:
continue
s_used.add(node)
for i in xrange(len(e)):
temp = node + e[i]
temp = cap(temp, k)
temp = tuple(temp)
nodes[temp] = 0
if not temp in s_used:
s.add(temp)
return nodes
def cap(t, k):
for i in xrange(len(t)):
if t[i] > k:
t[i] = k
return t
我不喜欢它。
字典nodes
的键是期望的向量。
答案 0 :(得分:2)
from itertools import product
def nodes_iter(n, k):
""" returns generator (lazy iterator) rather than creating whole list """
return product(range(k+1),repeat=n)
使用示例:
for node in nodes_iter(3,1):
print node
(0, 0, 0)
(0, 0, 1)
(0, 1, 0)
(0, 1, 1)
(1, 0, 0)
(1, 0, 1)
(1, 1, 0)
(1, 1, 1)