Josh
Josh

Reputation: 151

SageMath: Creating a list of vectors of an arbitrary dimension ranging from -a to +a

I would like to create a list of vectors of an arbitrary dimension n such that all of their entries range from an input -a to a. I know how to do this for a given dimension. For instance if n=2, the following does the job:

a = 1
[vector([x,y]) for x in [-a..a] for y in [-a..a]]

#output:
[(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 0), (0, 1), (1, -1), (1, 0), (1, 1)]

I am not sure how to extend this for an arbitrary n. I have used list comprehension, but it can be done in any other format. One idea that came to mind, which did not work was to use something like:

n = 2
letters = [chr(ord('`')+i+23) for i in [1..n]]

# output
# ['x', 'y']

a = 1
[vector([var(letters[0]),var(letters[1])]) for var(letters[0]) in [-a..a] for var(letters[1]) in [-a..a]]

# SyntaxError: can't assign to function call

Is there an efficient way of doing this in SageMath?

Upvotes: 1

Views: 579

Answers (1)

Will Da Silva
Will Da Silva

Reputation: 7040

You can use itertools.product for this:

import itertools as it


def all_coords(a: int, ndim: int):
    yield from it.product(range(-a + 1, a), repeat=ndim)

Usage:

>>> list(all_coords(2, ndim=2))
[(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 0), (0, 1), (1, -1), (1, 0), (1, 1)]
>>> list(all_coords(3, ndim=2))
[(-2, -2), (-2, -1), (-2, 0), (-2, 1), (-2, 2), (-1, -2), (-1, -1), (-1, 0), (-1, 1), (-1, 2), (0, -2), (0, -1), (0, 0), (0, 1), (0, 2), (1, -2), (1, -1), (1, 0), (1, 1), (1, 2), (2, -2), (2, -1), (2, 0), (2, 1), (2, 2)]
>>> list(all_coords(2, ndim=3))
[(-1, -1, -1), (-1, -1, 0), (-1, -1, 1), (-1, 0, -1), (-1, 0, 0), (-1, 0, 1), (-1, 1, -1), (-1, 1, 0), (-1, 1, 1), (0, -1, -1), (0, -1, 0), (0, -1, 1), (0, 0, -1), (0, 0, 0), (0, 0, 1), (0, 1, -1), (0, 1, 0), (0, 1, 1), (1, -1, -1), (1, -1, 0), (1, -1, 1), (1, 0, -1), (1, 0, 0), (1, 0, 1), (1, 1, -1), (1, 1, 0), (1, 1, 1)]

Upvotes: 1

Related Questions