Reputation: 151
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
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