同时迭代多个numpy数组

Question

I have a numpy array with shape (n,), for example:

[
    [0, 1],         # A
    [0, 1, 2, 3],   # B
    [0, 1, 2]       # C
]

How can I iterate through all 'tuples' formed by this array, in this case I would expect the return to be (0,0,0),(1,0,0),(0,1,0),(0,2,0),(0,3,0),(1,1,0)... and so on

Answer 1

Here is a numpy solution for you.

inp = [
    [0, 1,],        
    [0, 1, 2, 3],  
    [0, 1, 2]       
]

Now you can simply

import numpy as np

shp = []
for sub_list in inp:
    shp.append(len(sub_list))

arr = np.ones(shp)

result = np.where(arr)

tuples = [t for t in zip(*result)]

out:

 [(0, 0, 0), (0, 0, 1), (0, 0, 2), (0, 1, 0), (0, 1, 1), (0, 1, 2), (0, 2, 0), (0, 2, 1), (0, 2, 2), (0, 3, 0), (0, 3, 1), (0, 3, 2), (1, 0, 0), (1, 0, 1), (1, 0, 2), (1, 1, 0), (1, 1, 1), (1, 1, 2), (1, 2, 0), (1, 2, 1), (1, 2, 2), (1, 3, 0), (1, 3, 1), (1, 3, 2)]

The way this works, is you build an array of ones whose dimensions are the lengths of your lists.
Then you get the multi dimensional indices of this array, which happen to be exactly what you wanted.

If you also want to access your lists in the relevant index, you can easily do that as well.