python,创建矢量的优雅方式

时间:2012-04-22 19:27:52

标签: python

创建维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的键是期望的向量。

1 个答案:

答案 0 :(得分: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)