Most efficient way to create an (n*m by m) numpy array with the first row all zeros and last row all n in Python?

Question

I am looking for the most effective way to create an array that has m columns so that the final outcome is:

[[0, 0, 0, ..., 0]
[1, 0, 0, ..., 0]
[2, 0, 0, ..., 0]
[3, 0, 0, ..., 0]
[4, 0, 0, ..., 0]
[5, 0, 0, ..., 0]
...
[n, 0, 0, ..., 0]
[n, 1, 0, ..., 0]
[n, 2, 0, ..., 0]
...
[n, n, 0, ..., 0]
...
[n, n, n, ..., n-3]
[n, n, n, ..., n-2]
[n, n, n, ..., n-1]
[n, n, n, ..., n]]

If n = 2 and m = 3 we would have:

[[0, 0, 0]
[1, 0, 0]
[2, 0, 0]
[0, 1, 0]
[1, 1, 0]
[2, 1, 0]
[0, 2, 0]
[1, 2, 0]
[2, 2, 0]
[0, 0, 1]
[1, 0, 1]
[2, 0, 1]
[0, 1, 1]
[1, 1, 1]
[2, 1, 1]
[0, 2, 1]
[1, 2, 1]
[2, 2, 1]
[0, 0, 2]
[1, 0, 2]
[2, 0, 2]
[0, 1, 2]
[1, 1, 2]
[2, 1, 2]
[0, 2, 2]
[1, 2, 2]
[2, 2, 2]]

I am trying to avoid as much as possible for-loops, and I would like to be able to scale it up so that I can give any positive, non-null (n,m).

Bonus question: is there a name for such an array?

I have tried np.meshgrid(), but that does not yield to the right array.
I tried np.repeat(), and also failed.

Asked By: besp

||

Source

Answer 1

You can use the itertools product function and do:

n = 2
m = 3

# note: I do perm[::-1] to reverse the output of
# product to match the required order in the question
output = np.array(
    [perm[::-1] for perm in product(range(n + 1), repeat=m)]
)

print(output)

[[0 0 0]
 [1 0 0]
 [2 0 0]
 [0 1 0]
 [1 1 0]
 [2 1 0]
 [0 2 0]
 [1 2 0]
 [2 2 0]
 [0 0 1]
 [1 0 1]
 [2 0 1]
 [0 1 1]
 [1 1 1]
 [2 1 1]
 [0 2 1]
 [1 2 1]
 [2 2 1]
 [0 0 2]
 [1 0 2]
 [2 0 2]
 [0 1 2]
 [1 1 2]
 [2 1 2]
 [0 2 2]
 [1 2 2]
 [2 2 2]]

Answered By: Matt Pitkin

Most efficient way to create an (n*m by m) numpy array with the first row all zeros and last row all n in Python?

Question:

Answers: