2012-04-22 95 views
0

什么是优雅的方式来创建维度N的所有向量的集合,每个元素都是介于0和K之间的整数([0,K])。python,创建向量的优雅方式

我当前的代码是:

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

出于好奇,为什么你要这样?你不会以K^N向量结束吗? (很多) – Shep 2012-04-22 19:39:18

+0

我正在玩树木,在我的情况下,n和k不是这样,计算机可以处理它 – ashim 2012-04-22 19:50:01

回答

2

使用itertools

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) 
+0

'product'采用关键字参数'repeat'。所以,产品(范围(k + 1),重复= n)'会做。并注意'k'是包含性的。所以'range(k + 1)'而不是'range(k)'。 – Avaris 2012-04-22 19:40:21

+0

我冒昧地根据Avaris的评论编辑答案(并在控制台上测试之后)。 – 2012-04-22 19:42:53

+0

@Roberto:感谢编辑,但我在编辑 – vartec 2012-04-22 19:45:59