CODEWITHSUNDEEP
python
ashim
ashim

Reputation: 25570

python, elegant way of creating vectors

what is elegant way to create set of ALL vectors of dimension N, that each element is integer between 0 and K inclusive ([0, K]).

my current code is :

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

and I don't like it. keys of dictionary nodes are desired vectors.

Upvotes: 0

Views: 122

Answers (1)

vartec
vartec

Reputation: 134641

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

Example usage:

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)

Upvotes: 2

Related Questions