How to effectively generate an array of tuples using numpy

Question

I would like to effectively generate a numpy array of tuples which size is the multiple of the dimensions of each axis using numpy.arange() and exclusively using numpy functions. For example: the size of a_list below is max_i*max_j*max_k.

Moreover, the array that I would like to obtain for the example below looks like this : [(0,0,0), (0,0,1), ..., (0, 0, 9), (0, 1, 0), (0, 1, 1), ..., (9, 4, 14)]

a_list = list()
max_i = 10
max_j = 5
max_k = 15

for i in range(0, max_i):
  for j in range(0, max_j):
    for k in range(0, max_k):
      a_list.append((i, j, k))

The loop's complexity above, relying on list and for loops, is O(max_i*max_j*max_k), I would like to use a factorized way to generate a lookalike array of tuples in numpy. Is it possible ?

tuples are fundamentally a part of python, not numpy. Can you use arrays with a 3d dimension? Or going the other way, itertools? — Mad Physicist, Nov 06 '19 at 14:05
Technically, your operation is O(m * n * p), not O(n * n * n) since you have presumably three independent quantities. The triple loop has optimal complexity. You won't do better than linear time in each variable, even using numpy. — Mad Physicist, Nov 06 '19 at 14:07
Yes thank you for this last remark, I'm going to correct this — YellowishLight, Nov 06 '19 at 14:08
So what I'm saying is please clarify what you want. If you want to avoid using the extra space and just want an iterator, I would suggest itertools. If you want densely packed data, I would suggest an (M, N, P, 3) numpy array. Either way, please explain what you actually want, because right now it doesn't make much sense. — Mad Physicist, Nov 06 '19 at 14:09
You can get a 2D array with `np.indices((max_i,max_j,max_k)).transpose(1,2,3,0).reshape(-1,3)`. — Divakar, Nov 06 '19 at 14:12
Also, `k` changes fastest, so the second element will be `(0, 0, 1)`, not `(1, 0, 0)`. — Mad Physicist, Nov 06 '19 at 14:13

score 1 · Answer 1 · answered Nov 06 '19 at 15:18

I like Divakar's solution in the comments better, but here's another.

What you're describing is a cartesian product. With some help from this post, you can achieve this as follows

import numpy as np

# Input
max_i, max_j, max_k = (10, 5, 15)

# Build sequence arrays 0, 1, ... N
arr_i = np.arange(0, max_i)
arr_j = np.arange(0, max_j)
arr_k = np.arange(0, max_k)

# Build cartesian product of sequence arrays
grid = np.meshgrid(arr_i, arr_j, arr_k)
cartprod = np.stack(grid, axis=-1).reshape(-1, 3)

# Convert to list of tuples
result = list(map(tuple, cartprod))

How to effectively generate an array of tuples using numpy

1 Answers1